tico 0.1.0__py3-none-any.whl

tico/serialize/operators/op_pow.py

@@ -0,0 +1,141 @@
+# 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 extract_torch_dtype, to_circle_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 PowTensorScalarArgs, PowTensorTensorArgs
+
+
+class BasePowVisitor(NodeVisitor):
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def cast_to_float(self, node: torch.fx.Node) -> circle.Tensor.TensorT:
+        assert isinstance(node, torch.fx.Node), type(node)
+        node_tensor: circle.Tensor.TensorT = self.graph.get_tensor(node)
+        node_shape: List[int] = node_tensor.shape
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.CAST, self._op_codes
+        )
+        cast_name = f"{node.name}_cast"
+        cast_dtype = circle.TensorType.TensorType.FLOAT32
+        cast_tensor = self.graph.add_tensor_from_scratch(
+            prefix=cast_name,
+            dtype=cast_dtype,
+            shape=node_shape,
+            source_node=node,
+        )
+        cast_operator = create_builtin_operator(
+            self.graph, op_index, [node], [cast_tensor]
+        )
+        cast_operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.CastOptions
+        )
+        option = circle.CastOptions.CastOptionsT()
+        node_dtype = extract_torch_dtype(node)
+        option.inDataType = to_circle_dtype(node_dtype)
+        option.outDataType = cast_dtype
+        cast_operator.builtinOptions = option
+        self.graph.add_operator(cast_operator)
+
+        return cast_tensor
+
+    def define_pow_node(self, inputs: List, outputs: List) -> circle.Operator.OperatorT:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.POW, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.PowOptions
+        # Pow opearation does not have any options.
+        option = circle.PowOptions.PowOptionsT()
+
+        operator.builtinOptions = option
+
+        return operator
+
+
+# TODO Support `aten::pow.Scalar` (base=scalar, exponenent=tensor)
+# ExecuTorch currently does not support it as of now (2024/02/13).
+
+
+@register_node_visitor
+class PowTensorScalarVisitor(BasePowVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.pow.Tensor_Scalar]
+
+    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 = PowTensorScalarArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        exponent = args.exponent
+
+        lhs_dtype = extract_torch_dtype(input)
+        # Circle supports only same dtype between lhs and rhs.
+        if lhs_dtype == torch.float32 and isinstance(exponent, int):
+            exponent = float(exponent)
+        if lhs_dtype == torch.int32 or lhs_dtype == torch.int64:
+            if isinstance(exponent, float):
+                input = self.cast_to_float(input)  # type: ignore[assignment]
+
+        operator = self.define_pow_node([input, exponent], [node])
+
+        return operator
+
+
+@register_node_visitor
+class PowTensorTensorVisitor(BasePowVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.pow.Tensor_Tensor]
+
+    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 = PowTensorTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        exponent = args.exponent  # type: ignore[arg-type]
+
+        lhs_dtype = extract_torch_dtype(input)
+        rhs_dtype = extract_torch_dtype(exponent)  # type: ignore[arg-type]
+        # Circle supports only same dtype between lhs and rhs.
+        if lhs_dtype == torch.float32 and rhs_dtype == torch.int:
+            exponent = self.cast_to_float(exponent)  # type: ignore[arg-type, assignment]
+        if lhs_dtype == torch.int32 or lhs_dtype == torch.int64:
+            if rhs_dtype == torch.float32:
+                input = self.cast_to_float(input)  # type: ignore[assignment]
+
+        operator = self.define_pow_node([input, exponent], [node])
+
+        return operator

tico/serialize/operators/op_prelu.py

@@ -0,0 +1,54 @@
+# 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 PReLUArgs
+
+
+@register_node_visitor
+class PReLUVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.prelu.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:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.PRELU, self._op_codes
+        )
+
+        args = PReLUArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        weight = args.weight
+
+        inputs = [input, weight]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        return operator

tico/serialize/operators/op_quantize_per_tensor.py

@@ -0,0 +1,79 @@
+# 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 QuantizePerTensorArgs
+
+
+@register_node_visitor
+class QuantizePerTensorDefaultVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.quantized_decomposed.quantize_per_tensor.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 = QuantizePerTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        tensor = args.tensor
+        scale = args.scale
+        zero_p = args.zero_p
+        quant_min = args.quant_min
+        quant_max = args.quant_max
+
+        output_tensor: circle.Tensor.TensorT = self.graph.get_tensor(node)
+        assert output_tensor.quantization is not None
+
+        # Tensor should have qparam when it's exported
+        # The qparam must match with the arguments of this Op
+        assert output_tensor.quantization.scale[0] == scale
+        assert output_tensor.quantization.zeroPoint[0] == zero_p
+
+        if output_tensor.type == circle.TensorType.TensorType.UINT8:
+            assert quant_min == 0 and quant_max == 255
+        elif output_tensor.type == circle.TensorType.TensorType.INT16:
+            # Some frameworks use -32767 as quant_min of int16
+            assert quant_min in (-32768, -32767) and quant_max == 32767
+
+        inputs = [tensor]
+        outputs = [node]
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.QUANTIZE, self._op_codes
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.QuantizeOptions
+        )
+        option = circle.QuantizeOptions.QuantizeOptionsT()
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_reciprocal.py

@@ -0,0 +1,64 @@
+# 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 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 ReciprocalArgs
+
+
+@register_node_visitor
+class ReciprocalVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.reciprocal.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:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.DIV, self._op_codes
+        )
+
+        args = ReciprocalArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+
+        input_tensor = torch.tensor(1, dtype=extract_torch_dtype(input))
+        x = self.graph.add_const_tensor(input_tensor, source_node=node)
+        inputs = [x, input]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.DivOptions
+        option = circle.DivOptions.DivOptionsT()
+        option.fusedActivationFunction = (
+            circle.ActivationFunctionType.ActivationFunctionType.NONE
+        )
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_relu.py

@@ -0,0 +1,53 @@
+# 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 ReluArgs
+
+
+@register_node_visitor
+class ReluVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.relu.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:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.RELU, self._op_codes
+        )
+
+        args = ReluArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+
+        inputs = [input]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        return operator

tico/serialize/operators/op_relu6.py

@@ -0,0 +1,52 @@
+# 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 Relu6Args
+
+
+@register_node_visitor
+class Relu6Visitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.relu6.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 = Relu6Args(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.RELU6, self._op_codes
+        )
+
+        inputs = [input]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        return operator

tico/serialize/operators/op_repeat.py

@@ -0,0 +1,100 @@
+# 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 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.errors import InvalidArgumentError, NotYetSupportedError
+from tico.utils.validate_args_kwargs import RepeatArgs
+
+
+@register_node_visitor
+class RepeatVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.repeat.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 = RepeatArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        repeats = args.repeats
+
+        for r in repeats:
+            if r == 0:
+                # TODO: Support r == 0 case
+                raise NotYetSupportedError("Only support positive repeat value")
+            elif r < 0:
+                raise InvalidArgumentError("Only support positive repeat value")
+
+        tensor_shape = extract_shape(input)
+        assert len(tensor_shape) <= len(repeats)
+        if len(tensor_shape) != len(repeats):
+            # TODO Support len(tensor_shape) < len(repeats)
+            raise NotYetSupportedError(
+                "Length of both input tensor and repeats vector should be same."
+            )
+        repeat_dim_cnt = len(repeats) - repeats.count(1)
+        tensor_dtype = extract_circle_dtype(input)
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.CONCATENATION, self._op_codes
+        )
+        concat_input: torch.fx.Node | circle.Tensor.TensorT = input
+        concat_output: torch.fx.node.Node | circle.Tensor.TensorT = node
+        for idx, r in enumerate(repeats):
+            # concat along idx dimension
+            if r > 1:
+                # Except last created concat, a tensor should be created.
+                if repeat_dim_cnt > 1:
+                    repeated_shape = list(tensor_shape)
+                    repeated_shape[idx] = repeated_shape[idx] * r
+                    concat_output = self.graph.add_tensor_from_scratch(
+                        prefix=f"{node.name}_concat_{idx}",
+                        shape=repeated_shape,
+                        dtype=tensor_dtype,
+                        source_node=node,
+                    )
+                inputs = [concat_input] * r
+                if repeat_dim_cnt == 1:
+                    outputs: List[torch.fx.node.Node | circle.Tensor.TensorT] = [node]
+                else:
+                    outputs = [concat_output]
+                operator = create_builtin_operator(
+                    self.graph, op_index, inputs, outputs
+                )
+                operator.builtinOptionsType = (
+                    circle.BuiltinOptions.BuiltinOptions.ConcatenationOptions
+                )
+                option = circle.ConcatenationOptions.ConcatenationOptionsT()
+                option.axis = idx
+                operator.builtinOptions = option
+                if repeat_dim_cnt > 1:
+                    self.graph.add_operator(operator)
+                    concat_input = concat_output
+                repeat_dim_cnt -= 1
+
+        return operator

tico/serialize/operators/op_reshape.py

@@ -0,0 +1,73 @@
+# 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, is_const
+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.errors import NotYetSupportedError
+from tico.utils.validate_args_kwargs import ReshapeArgs
+
+
+@register_node_visitor
+class ReshapeVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.reshape.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:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.RESHAPE,
+            self._op_codes,
+        )
+        args = ReshapeArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        size = args.shape
+
+        if isinstance(size, int):
+            raise NotYetSupportedError("scalar size conversion is not supported yet.")
+
+        assert is_const(size), type(size)
+
+        size_i32 = circle_legalize_dtype_to(size, dtype=torch.int32)
+        inputs = [input, size_i32]
+        outputs = [node]
+
+        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 = size_i32
+
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_resize_nearest_neighbor.py

@@ -0,0 +1,70 @@
+# 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 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 ResizeNearestNeighborArgs
+
+
+@register_node_visitor
+class ResizeNearestNeighborVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.circle_custom.resize_nearest_neighbor
+    ]
+
+    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:
+        # TODO Support generic algorithm
+        args = ResizeNearestNeighborArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        size = args.size
+
+        size_i32 = circle_legalize_dtype_to(size, dtype=torch.int32)
+        inputs = [input, size_i32]
+        outputs = [node]
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.RESIZE_NEAREST_NEIGHBOR,
+            self._op_codes,
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.ResizeNearestNeighborOptions
+        )
+        option = circle.ResizeNearestNeighborOptions.ResizeNearestNeighborOptionsT()
+        # TODO Consider these options
+        # If True, the centers of the 4 corner pixels of the input and output tensors are aligned, preserving the values at the corner pixels.
+        option.alignCorners = False
+        # If True, the pixel centers are assumed to be at (0.5, 0.5). If this parameter is True, then align_corners parameter must be False.
+        option.halfPixelCenters = False
+        operator.builtinOptions = option
+
+        return operator