tico 0.1.0__py3-none-any.whl

tico/serialize/operators/op_clamp.py

@@ -0,0 +1,126 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.passes import ops
+
+from tico.serialize.circle_graph import (
+    CircleSubgraph,
+    extract_circle_dtype,
+    extract_shape,
+)
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import ClampArgs
+
+
+@register_node_visitor
+class ClampVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = ops.aten.clamp
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_minimum_node(
+        self,
+        inputs: List[torch.fx.Node | circle.Tensor.TensorT | int | float],
+        outputs: List[torch.fx.Node | circle.Tensor.TensorT],
+    ) -> circle.Operator.OperatorT:
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.MINIMUM, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.MaximumMinimumOptions
+        )
+        option = circle.MaximumMinimumOptions.MaximumMinimumOptionsT()
+
+        operator.builtinOptions = option
+        return operator
+
+    def define_maximum_node(
+        self,
+        inputs: List[torch.fx.Node | circle.Tensor.TensorT | int | float],
+        outputs: List[torch.fx.Node | circle.Tensor.TensorT],
+    ) -> circle.Operator.OperatorT:
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.MAXIMUM, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.MaximumMinimumOptions
+        )
+        option = circle.MaximumMinimumOptions.MaximumMinimumOptionsT()
+
+        operator.builtinOptions = option
+
+        return operator
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = ClampArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        min_val = args.min
+        max_val = args.max
+
+        if min_val is None and max_val is None:
+            raise ValueError("Both min and max cannot be None")
+
+        elif min_val is not None and max_val is None:
+            # min only
+            return self.define_maximum_node([input, min_val], [node])
+
+        elif min_val is None and max_val is not None:
+            # max only
+            return self.define_minimum_node([input, max_val], [node])
+
+        elif min_val is not None and max_val is not None:
+            input_shape = extract_shape(input)
+            input_dtype = extract_circle_dtype(input)
+            minimum_tensor = self.graph.add_tensor_from_scratch(
+                prefix=f"{input.name}_min",
+                dtype=input_dtype,
+                shape=list(input_shape),
+                source_node=node,
+            )
+            minimum_opertor = self.define_minimum_node(
+                [input, max_val], [minimum_tensor]
+            )
+            self.graph.add_operator(minimum_opertor)
+
+            maximum_operator = self.define_maximum_node(
+                [minimum_tensor, min_val], [node]
+            )
+            return maximum_operator
+
+        else:
+            raise RuntimeError("Cannot reach here")

tico/serialize/operators/op_clone.py

@@ -0,0 +1,71 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import CloneArgs
+
+
+@register_node_visitor
+class CloneVisitor(NodeVisitor):
+    """
+    Clone tensor
+    TODO: Support dim_order and memory_format
+    Tranpose may be required if 'memory_format' differs from input tensor's 'memory_format'
+    """
+
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.clone.default]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        if "memory_format" in node.kwargs:
+            # TODO: Support dim_order and memory_format
+            pass
+
+        args = CloneArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.TRANSPOSE, self._op_codes
+        )
+
+        permute = torch.IntTensor(list(range(len(input.meta["val"].shape))))
+
+        inputs = [input, permute]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.TransposeOptions
+        )
+        option = circle.TransposeOptions.TransposeOptionsT()
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_constant_pad_nd.py

@@ -0,0 +1,72 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_mapping import extract_shape
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.errors import InvalidArgumentError
+from tico.utils.validate_args_kwargs import ConstantPadNdArgs
+
+
+@register_node_visitor
+class ConstantPadNdVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.constant_pad_nd.default]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = ConstantPadNdArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input_ = args.input
+        pad = args.pad
+        val = args.value
+
+        if val != 0:
+            raise InvalidArgumentError("Only support 0 value padding.")
+
+        input_shape_len = len(extract_shape(input_))
+        padding_size = [[pad[2], pad[3]], [pad[0], pad[1]]]
+        if input_shape_len == 3:
+            padding_size = [[0, 0]] + padding_size
+        elif input_shape_len == 4:
+            padding_size = [[0, 0], [0, 0]] + padding_size
+        else:
+            raise InvalidArgumentError("Only support 3D/4D inputs.")
+
+        paddings = torch.tensor(padding_size, dtype=torch.int32)
+        inputs = [input_, paddings]
+        outputs = [node]
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.PAD, self._op_codes
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.PadOptions
+        option = circle.PadOptions.PadOptionsT()
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_conv2d.py

@@ -0,0 +1,186 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_mapping import extract_circle_dtype, extract_shape
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.serialize.quant_param import QPARAM_KEY, QuantParam
+from tico.utils.define import define_pad_node
+from tico.utils.padding import is_same_padding, is_valid_padding, SAME, VALID
+from tico.utils.validate_args_kwargs import Conv2DArgs
+
+
+@register_node_visitor
+class Conv2dVisitor(NodeVisitor):
+    """
+    NOTE
+    - The padding of CircleConv2D has only padding type('VALID', 'SAME') in circle, but the padding of nn.Conv2d has padding type(('valid', 'same')), padding value(int)
+    and padding value(tuple->[pad_h, pad_w]).
+    ref: https://tensorflow.org/api_docs/python/tf/nn/conv2d
+
+    [1] With valid/same padding: CircleConv2D (only)
+
+        [ATEN IR]
+        Input[NHWC] ---- circle_cumstom.conv2d[NHWC] ---- OUTPUT[NHWC]
+        Weight[NHWC] ---/
+        Bias ----------/
+
+        [CIRCLE IR]
+        Input[NHWC] ----  CircleConv2D[NHWC] ---- OUTPUT[NHWC]
+        Weight[NHWC] ---/
+        Bias ----------/
+
+    [2] With additional padding: CirclePad + CircleConv2D
+
+        [ATEN IR]
+        Input[NHWC] ---- circle_cumstom.conv2d[NHWC] ---- OUTPUT[NHWC]
+        Weight[NHWC] ---/
+        Bias ----------/
+
+        [CIRCLE IR]
+        Input[NHWC] ---- CirclePad[NHWC] ---- CircleConv2D[NHWC] ---- OUTPUT[NHWC]
+                         Weight[NHWC] ------/
+                         Bias -------------/
+    """
+
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.circle_custom.conv2d,
+        torch.ops.circle_custom.conv2d.padding,
+    ]
+
+    def define_conv2d_node(
+        self, padding: int, stride: List, dilation: List, inputs: List, outputs: List
+    ) -> circle.Operator.OperatorT:
+        def set_conv2d_option(operator, stride, dilation):
+            operator.builtinOptionsType = (
+                circle.BuiltinOptions.BuiltinOptions.Conv2DOptions
+            )
+            option = circle.Conv2DOptions.Conv2DOptionsT()
+            option.padding = padding
+            option.strideH = stride[0]
+            option.strideW = stride[1]
+            option.dilationHFactor = dilation[0]
+            option.dilationWFactor = dilation[1]
+            option.fusedActivationFunction = (
+                circle.ActivationFunctionType.ActivationFunctionType.NONE
+            )
+            operator.builtinOptions = option
+
+        conv2d_op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.CONV_2D, self._op_codes
+        )
+        operator = create_builtin_operator(self.graph, conv2d_op_index, inputs, outputs)
+        set_conv2d_option(operator, stride, dilation)
+        return operator
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph):
+        super().__init__(op_codes, graph)
+
+    def define_node(self, node: torch.fx.Node) -> circle.Operator.OperatorT:
+        # conv2d(Tensor input, Tensor weight, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, SymInt[2] dilation=1, SymInt groups=1) -> Tensor
+        # conv2d.padding(Tensor input, Tensor weight, Tensor? bias=None, SymInt[2] stride=1, str padding="valid", SymInt[2] dilation=1, SymInt groups=1) -> Tensor
+        args = Conv2DArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+
+        input_ = args.input
+        weight = args.weight
+        bias = args.bias
+        stride = args.stride
+        padding = args.padding
+        dilation = args.dilation
+        groups = args.groups
+
+        assert groups == 1, "Only support group 1 conv2d"
+
+        input_dtype: int = extract_circle_dtype(input_)
+        input_shape = list(extract_shape(input_))
+        assert len(input_shape) == 4, len(input_shape)
+        output_shape = extract_shape(node)
+        assert len(output_shape) == 4, len(output_shape)
+
+        conv_input: torch.fx.node.Node | circle.Tensor.TensorT = input_
+        weight_shape = list(extract_shape(weight))
+
+        if is_valid_padding(padding):
+            conv2d_padding_type = VALID
+        elif is_same_padding(padding, input_shape, output_shape):
+            conv2d_padding_type = SAME
+        else:
+            assert isinstance(padding, list) and len(padding) == 2
+
+            conv2d_padding_type = VALID
+
+            # Padding is not valid or same, so we use valid padding and add padding operator before conv2d operator.
+            # when data_foramt is "NHWC", padding should be [[0, 0], [pad_top, pad_bottom], [pad_left, pad_right], [0, 0]]
+            paddings = torch.tensor(
+                [
+                    [0, 0],
+                    [padding[0], padding[0]],
+                    [padding[1], padding[1]],
+                    [0, 0],
+                ],
+                dtype=torch.int32,
+            )
+            pad_output_shape = [
+                input_shape[0],
+                input_shape[1],
+                input_shape[2],
+                input_shape[3],
+            ]
+            # Add (pad_top+pad_bottom) to pad_output_shape_h
+            pad_output_shape[1] += padding[0] * 2
+            # Add (pad_left+pad_Right) to pad_output_shape_w
+            pad_output_shape[2] += padding[1] * 2
+            # create padded output tensor
+            input_qparam: Optional[QuantParam] = (
+                input_.meta[QPARAM_KEY] if QPARAM_KEY in input_.meta else None
+            )
+            pad_output = self.graph.add_tensor_from_scratch(
+                prefix=f"{node.name}_input_pad_output",
+                shape=pad_output_shape,
+                dtype=input_dtype,
+                qparam=input_qparam,
+                source_node=node,
+            )
+            # CirclePad
+            pad_operator = define_pad_node(
+                self.graph, self._op_codes, [input_, paddings], [pad_output]
+            )
+            self.graph.add_operator(pad_operator)
+            conv_input = pad_output
+
+        if bias is None:
+            # luci-interpreter can't run no bias conv. Let's add zero vector for bias.
+            assert len(weight_shape) == 4
+            out_channel = weight_shape[0]
+            bias = [0.0] * out_channel  # type: ignore[assignment]
+
+        # Conv2D
+        conv2d_operator = self.define_conv2d_node(
+            conv2d_padding_type,  # 'SAME'(0) or 'VALID'(1)
+            stride,
+            dilation,
+            [conv_input, weight, bias],
+            [node],
+        )
+
+        return conv2d_operator

tico/serialize/operators/op_copy.py

@@ -0,0 +1,164 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING, Union
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.errors import NotYetSupportedError
+from tico.utils.validate_args_kwargs import CopyArgs
+
+
+@register_node_visitor
+class CopyVisitor(NodeVisitor):
+    """
+    NOTE `torch.Tensor.copy_`'s behavior matches with `Reshape` of CIRCLE.
+    - because `torch.Tensor.copy_` is a in-place operator, so `dst` is converted to `Shape` of CIRCLE.
+    - after that, `dst` converted to `Shape` is connected to shape of `Reshape`.
+    - `src` is connected to tensor of `Reshape`.
+    - if `dst` is not converted to `Shape`.
+      [dst]      [src]
+                   |
+                [Reshape]
+    - if `dst` is converted to `Shape`.
+      [dst]      [src]
+        |          |
+      [Shape]      |
+        \         /
+         [Reshape]
+    """
+
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.copy.default]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def check_to_do_broadcast(self, dst: List[int], src: List[int]) -> bool:
+        return dst != src
+
+    def define_broadcast_to_node(
+        self,
+        inputs: List[Union[circle.Tensor.TensorT, torch.Tensor]],
+        outputs: List[circle.Tensor.TensorT],
+    ) -> circle.Operator.OperatorT:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.BROADCAST_TO, self._op_codes
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.BroadcastToOptions
+        )
+
+        option = circle.BroadcastToOptions.BroadcastToOptionsT()
+        operator.builtinOptions = option
+        return operator
+
+    def define_shape_node(
+        self, inputs: List[torch.fx.Node], outputs: List[circle.Tensor.TensorT]
+    ) -> circle.Operator.OperatorT:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.SHAPE, self._op_codes
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.ShapeOptions
+
+        option = circle.ShapeOptions.ShapeOptionsT()
+        option.outType = circle.TensorType.TensorType.INT32
+        operator.builtinOptions = option
+        return operator
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        if len(node.args) == 3:
+            raise NotYetSupportedError("'non_blocking' is not supported yet.")
+
+        assert len(node.args) == 2, len(node.args)
+
+        args = CopyArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        dst = args.dst
+        src = args.src
+
+        # To connect 'dst' to Reshape node in the graph, 'dst' must be converted to Shape op.
+        dst_tensor: circle.Tensor.TensorT = self.graph.get_tensor(dst)
+        dst_shape: List[int] = dst_tensor.shape
+        dst_shape_tensor = torch.as_tensor(dst_shape, dtype=torch.int32)
+
+        dst_shape_shape = [len(dst_shape)]
+        dst_name: str = dst.name
+
+        shape_output = self.graph.add_tensor_from_scratch(
+            prefix=f"{dst_name}_shape_output",
+            shape=dst_shape_shape,
+            dtype=circle.TensorType.TensorType.INT32,
+            source_node=node,
+        )
+
+        shape_operator = self.define_shape_node([dst], [shape_output])
+        self.graph.add_operator(shape_operator)
+
+        src_tensor: circle.Tensor.TensorT = self.graph.get_tensor(src)
+        src_shape: List[int] = src_tensor.shape
+
+        # The src tensor must be broadcastable with the dst tensor.
+        do_broadcast = self.check_to_do_broadcast(dst_shape, src_shape)
+        if do_broadcast:
+            # create braodcastTo output tensor
+            src_name: str = src.name
+            src_type: int = src_tensor.type
+
+            broadcast_to_output: circle.Tensor.TensorT = (
+                self.graph.add_tensor_from_scratch(
+                    prefix=f"{src_name}_broadcast_to_output",
+                    shape=dst_shape,
+                    dtype=src_type,
+                    source_node=node,
+                )
+            )
+
+            broadcast_to_operator: circle.Operator.OperatorT = (
+                self.define_broadcast_to_node(
+                    [src_tensor, dst_shape_tensor], [broadcast_to_output]
+                )
+            )
+            self.graph.add_operator(broadcast_to_operator)
+            inputs: List = [broadcast_to_output, shape_output]
+        else:
+            inputs = [src, shape_output]
+
+        outputs = [node]
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.RESHAPE, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.ReshapeOptions
+        )
+        option = circle.ReshapeOptions.ReshapeOptionsT()
+        option.newShape = dst_shape
+
+        operator.builtinOptions = option
+        return operator

tico/serialize/operators/op_cos.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2025 Samsung Electronics Co., Ltd. All Rights Reserved
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+from tico.utils.validate_args_kwargs import CosArgs
+
+
+@register_node_visitor
+class CosVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.cos.default]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = CosArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.COS, self._op_codes
+        )
+
+        inputs = [input]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.CosOptions
+        option = circle.CosOptions.CosOptionsT()
+
+        operator.builtinOptions = option
+
+        return operator