tico 0.1.0__py3-none-any.whl

tico/serialize/operators/op_any.py

@@ -0,0 +1,150 @@
+# 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.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.circle_mapping import (
+    circle_legalize_dtype_to,
+    extract_circle_dtype,
+    extract_shape,
+    extract_torch_dtype,
+)
+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 AnyArgs
+
+
+@register_node_visitor
+class AnyVisitor(NodeVisitor):
+    """
+    Let's take NotEqual0 -> ReduceMax workaround for float, int
+    [RESTRICTION]
+        1. ReduceAny is not supported (luci-interpreter)
+    [CASE: BOOL]
+        (Bool tensors don't need 'Not Equal 0' at the first step.)
+        bool[d0..dN]      --- Reduce Max      ---> bool[]
+    [CASE: FLOAT, INT]
+        int/float[d0..dN] --- Not Equal 0     ---> bool[d0,...dN]
+                          --- Reduce Max      ---> bool[]
+        * [d0..dN] means a tensor with any shape
+        * [] means Scalar
+    """
+
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.any.default,
+        torch.ops.aten.any.dim,
+        torch.ops.aten.any.dims,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_max_node(
+        self, inputs: List, outputs: List, keepdims: bool
+    ) -> circle.Operator.OperatorT:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.REDUCE_MAX, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.ReducerOptions
+        )
+        option = circle.ReducerOptions.ReducerOptionsT()
+        option.keepDims = keepdims
+
+        operator.builtinOptions = option
+
+        return operator
+
+    def define_ne_node(self, inputs: List, outputs: List) -> circle.Operator.OperatorT:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.NOT_EQUAL, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.NotEqualOptions
+        )
+        option = circle.NotEqualOptions.NotEqualOptionsT()
+        operator.builtinOptions = option
+        return operator
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        args = AnyArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        dim = args.dim
+        keepdim = args.keepdim
+
+        input_shape = list(extract_shape(input))
+        output_shape = list(extract_shape(node))
+
+        dim_i32 = None
+        if dim is None:
+            dims = tuple(i for i in range(0, len(input_shape)))
+            dim_i32 = tuple(
+                circle_legalize_dtype_to(dim, dtype=torch.int32) for dim in dims
+            )
+        if isinstance(dim, int):
+            dim_i32 = circle_legalize_dtype_to(dim, dtype=torch.int32)
+        if isinstance(dim, tuple):
+            dim_i32 = tuple(circle_legalize_dtype_to(d, dtype=torch.int32) for d in dim)
+        assert dim_i32 is not None
+
+        inputs = [
+            input,
+            dim_i32,
+        ]  # type: ignore[list-item]
+        outputs = [node]
+
+        dtype_torch = extract_torch_dtype(input)
+        input_tensor: torch.fx.node.Node | circle.Tensor.TensorT = input
+
+        if dtype_torch in [torch.int32, torch.int64, torch.float32, torch.float64]:
+            dst_dtype_circle = circle.TensorType.TensorType.BOOL
+            dst_dtype_torch = torch.bool
+            ne_tensor: circle.Tensor.TensorT = self.graph.add_tensor_from_scratch(
+                prefix=f"{input.name}_ne",
+                shape=input_shape,
+                dtype=dst_dtype_circle,
+                source_node=input,
+            )
+            ne_node = self.define_ne_node(
+                [input_tensor, torch.Tensor([0]).to(dtype_torch)], [ne_tensor]
+            )
+            self.graph.add_operator(ne_node)
+
+            dtype_torch = dst_dtype_torch
+            input_tensor = ne_tensor
+            inputs = [ne_tensor, dim_i32]
+
+        inputs = [input_tensor, dim_i32]
+
+        reduce_node: circle.Operator.OperatorT = self.define_max_node(
+            inputs, outputs, keepdim
+        )
+
+        return reduce_node

tico/serialize/operators/op_arange_start_step.py

@@ -0,0 +1,61 @@
+# 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.utils.validate_args_kwargs import ArangeStartStepArgs
+
+
+@register_node_visitor
+class ArangeStartStepVisitor(NodeVisitor):
+    """
+    Fuse arange_start_step to const_tensor
+    """
+
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.arange.start_step]
+
+    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 = ArangeStartStepArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        start = args.start
+        end = args.end
+        step = args.step
+        delta = 1
+
+        if step is not None:
+            delta = step[0]  # type: ignore[index]
+        # assert False, "This pass must not be in use."
+
+        arange_dtype: torch.dtype = torch.float32
+        if isinstance(start, int) and isinstance(end, int):
+            arange_dtype = torch.int64
+
+        output_data = torch.arange(start=start, end=end, step=delta, dtype=arange_dtype)
+        self.graph.update_tensor_buffer(output_data, node.name)
+
+        return None  # type: ignore[return-value]

tico/serialize/operators/op_argmax.py

@@ -0,0 +1,62 @@
+# 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.circle_mapping import circle_legalize_dtype_to
+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 ArgMaxArgs
+
+
+@register_node_visitor
+class ArgMaxVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.argmax.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 = ArgMaxArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        tensor = args.tensor
+        dim = args.dim
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.ARG_MAX, self._op_codes
+        )
+
+        dim_i32 = circle_legalize_dtype_to(dim, dtype=torch.int32)
+        inputs = [tensor, dim_i32]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.ArgMaxOptions
+        option = circle.ArgMaxOptions.ArgMaxOptionsT()
+        option.outputType = circle.TensorType.TensorType.INT64
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_avg_pool2d.py

@@ -0,0 +1,192 @@
+# 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.
+
+import math
+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.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.utils.define import define_pad_node
+from tico.utils.errors import NotYetSupportedError
+from tico.utils.validate_args_kwargs import AvgPool2dArgs
+
+
+@register_node_visitor
+class AvgPool2DVisitor(NodeVisitor):
+    """
+    This class defines how to serialize AvgPool2D operation into Circle IR.
+
+    Torch                                           | Circle
+
+    count_include_pad: True/False                   | (count_include_pad): Always False
+    padding: number (could be valid, same, or etc)  | padding: "valid"/"same"
+
+    * Circle's avgpool2d has no option for count_include_pad, so we always set it as False.
+    """
+
+    target: List[torch._ops.OpOverload] = [torch.ops.circle_custom.avgpool2d]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def has_padding(self, args: AvgPool2dArgs) -> bool:
+        padding = args.padding
+        if padding[0] == 0 and padding[1] == 0:
+            return False
+        else:
+            return True
+
+    def has_same_padding(self, args: AvgPool2dArgs) -> bool:
+        input_shape = list(extract_shape(args.input))
+        kernel_size = args.kernel_size
+        stride = args.stride
+        assert stride
+        padding = args.padding
+        # TODO Update this function when supporting ceil_mode = True
+        assert args.ceil_mode is False
+        output_height = math.floor(
+            (input_shape[1] + padding[0] * 2 - kernel_size[0]) / stride[0] + 1
+        )
+        output_width = math.floor(
+            (input_shape[2] + padding[1] * 2 - kernel_size[1]) / stride[1] + 1
+        )
+
+        return input_shape[1] == output_height and input_shape[2] == output_width
+
+    def define_avgpool_node(self, inputs, outputs, padding, stride, kernel_size):
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.AVERAGE_POOL_2D,
+            self._op_codes,
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.Pool2DOptions
+        option = circle.Pool2DOptions.Pool2DOptionsT()
+
+        assert padding in {"SAME": 0, "VALID": 1}
+
+        option.padding = {"SAME": 0, "VALID": 1}[padding]
+        option.strideH = stride[0]
+        option.strideW = stride[1]
+        option.filterHeight = kernel_size[0]
+        option.filterWidth = kernel_size[1]
+        option.fusedActivationFunction = (
+            circle.ActivationFunctionType.ActivationFunctionType.NONE
+        )
+
+        operator.builtinOptions = option
+        return operator
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        """
+        PSEUDO CODE
+
+        if count_include_pad == True:
+            (Circle cannot represent count_include_pad=True in AvgPool2D. Therefore we manually add zero padding node.)
+            DEFINE zero padding node
+            DEFINE avgpool node with no padding (valid)
+        if count_include_pad == False:
+            (Lucky! Circle can represent count_include_pad=False)
+            DEFINE avgpool node with same/valid padding.
+
+            (However, it cannot represent all paddings. So, if the padding is not same or valid, we throw an error.)
+            if the paddding is neither same nor valid:
+                THROW an error.
+        """
+        args = AvgPool2dArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        kernel_size = args.kernel_size
+        stride = args.stride
+        padding = args.padding
+        count_include_pad = args.count_include_pad
+
+        avgpool_input: torch.fx.Node | circle.Tensor.TensorT = input
+
+        def define_padding_node():
+            assert isinstance(padding, list), type(padding)
+            padding_vec = torch.tensor(
+                [
+                    [0, 0],
+                    [padding[0], padding[0]],
+                    [padding[1], padding[1]],
+                    [0, 0],
+                ],
+                dtype=torch.int32,
+            )
+            input_shape = list(extract_shape(input))
+            input_dtype: int = extract_circle_dtype(input)
+            padded_input_shape = [
+                input_shape[0],
+                input_shape[1],
+                input_shape[2],
+                input_shape[3],
+            ]
+            padded_input_shape[1] += padding[0] * 2
+            padded_input_shape[2] += padding[1] * 2
+            # create padded input tensor
+            padded_input_tensor = self.graph.add_tensor_from_scratch(
+                prefix=f"{input.name}_pad_output",
+                shape=padded_input_shape,
+                dtype=input_dtype,
+                source_node=node,
+            )
+            pad_operator = define_pad_node(
+                self.graph, self._op_codes, [input, padding_vec], [padded_input_tensor]
+            )
+            self.graph.add_operator(pad_operator)
+            return padded_input_tensor
+
+        if count_include_pad is True:
+            # Add padding before avgpool2d
+            # Circle's avgpool2d does not support count_include_pad=True, so we need to add padding manually
+            if self.has_padding(args):
+                avgpool_input = define_padding_node()
+
+            result = self.define_avgpool_node(
+                [avgpool_input], [node], "VALID", stride, kernel_size
+            )
+        elif count_include_pad is False:
+            if not self.has_padding(args):  # valid padding
+                result = self.define_avgpool_node(
+                    [avgpool_input], [node], "VALID", stride, kernel_size
+                )
+            elif self.has_same_padding(args):
+                result = self.define_avgpool_node(
+                    [avgpool_input], [node], "SAME", stride, kernel_size
+                )
+            else:
+                # CASE: count_include_pad is False and not VALID/SAME padding
+                #
+                # Implement this when it's needed.
+                # If needed, may it help: the idea of ratio masking in https://github.com/Samsung/TICO/pull/119
+                raise NotYetSupportedError(
+                    f"Padding({padding}) with count_include_pad({count_include_pad}) is not supported yet."
+                )
+        else:
+            raise RuntimeError("Cannot reach here")
+
+        return result

tico/serialize/operators/op_bmm.py

@@ -0,0 +1,62 @@
+# 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 BmmArgs
+
+
+@register_node_visitor
+class BatchMatmulVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.bmm.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 = BmmArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        mat2 = args.mat2
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.BATCH_MATMUL, self._op_codes
+        )
+
+        inputs = [input, mat2]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.BatchMatMulOptions
+        )
+        option = circle.BatchMatMulOptions.BatchMatMulOptionsT()
+        option.adjointLhs, option.adjointRhs = False, False
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_cat.py

@@ -0,0 +1,66 @@
+# 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 CatArgs
+
+
+@register_node_visitor
+class CatVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.cat.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 = CatArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        tensors = args.tensors
+        dim = args.dim
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.CONCATENATION, self._op_codes
+        )
+
+        inputs = tensors
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.ConcatenationOptions
+        )
+        option = circle.ConcatenationOptions.ConcatenationOptionsT()
+
+        option.axis = dim
+
+        option.fusedActivationFunction = (
+            circle.ActivationFunctionType.ActivationFunctionType.NONE
+        )
+        operator.builtinOptions = option
+
+        return operator