tico 0.1.0.dev250411__py3-none-any.whl

tico/serialize/operators/op_sub.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.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 SubTensorArgs
+
+
+@register_node_visitor
+class SubVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.sub.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:
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.SUB, self._op_codes
+        )
+
+        args = SubTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type, call-arg]
+
+        input = args.input
+        other = args.other
+        alpha = args.alpha
+
+        if alpha is not None:
+            raise NotYetSupportedError(
+                "'alpha' of aten::sub.Tensor is not supported yet"
+            )
+
+        inputs = [input, other]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.SubOptions
+        option = circle.SubOptions.SubOptionsT()
+        option.fusedActivationFunction = (
+            circle.ActivationFunctionType.ActivationFunctionType.NONE
+        )
+        option.potScaleInt16 = False
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_sum.py

@@ -0,0 +1,63 @@
+# 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 SumDimIntListArgs
+
+
+@register_node_visitor
+class SumVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.sum.dim_IntList]
+
+    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 = SumDimIntListArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        dim = args.dim
+        keepdim = args.keepdim
+
+        dim_i32 = circle_legalize_dtype_to(dim, dtype=torch.int32)
+
+        inputs = [input, dim_i32]
+        outputs = [node]
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.SUM, 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 = keepdim
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_tanh.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 TanhArgs
+
+
+@register_node_visitor
+class TanhVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.tanh.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 = TanhArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+
+        inputs = [input]
+        outputs = [node]
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.TANH, self._op_codes
+        )
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        return operator

tico/serialize/operators/op_to_copy.py

@@ -0,0 +1,105 @@
+# 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_circle_dtype,
+    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.errors import NotYetSupportedError
+from tico.utils.validate_args_kwargs import ToCopyArgs
+
+
+@register_node_visitor
+class ToCopyVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten._to_copy.default]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph):
+        super().__init__(op_codes, graph)
+
+    def define_cast_node(
+        self,
+        inputs: List[torch.fx.Node],
+        outputs: List[torch.fx.Node],
+        in_type: int,
+        out_type: int,
+    ):
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.CAST, self._op_codes
+        )
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.CastOptions
+        option = circle.CastOptions.CastOptionsT()
+        option.inDataType = in_type
+        option.outDataType = out_type
+        operator.builtinOptions = option
+
+        return operator
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        supported_kwargs = ["dtype", "device", "layout"]
+        if not all(k in supported_kwargs for k in node.kwargs):
+            unsupported_node_kargs = list(node.kwargs.keys())
+            for supported_key in supported_kwargs:
+                if supported_key in node.kwargs:
+                    unsupported_node_kargs.remove(supported_key)
+            raise NotYetSupportedError(
+                f"Support only {supported_kwargs} kwargs now. Do not support {unsupported_node_kargs}"
+            )
+
+        args = ToCopyArgs(*node.args, **node.kwargs)  # type: ignore[arg-type, call-arg]
+        input = args.input
+        dtype = args.dtype
+
+        input_meta = input.meta["val"]
+        # https://pytorch.org/docs/stable/tensor_attributes.html#torch-layout
+        # layout is two types: torch.strided(dense Tensors), torch.sparse_coo(sparse COO Tensors)
+        if "layout" in input.kwargs and input.kwargs["layout"] != input_meta:
+            raise NotYetSupportedError(
+                f"Only support when node and its input have same layout: (input layout: {input_meta}), (node layout: {node.kwargs['layout']})."
+            )
+
+        if dtype is not None:
+            target_type = node.kwargs["dtype"]
+        else:
+            # device and layout are meaningless
+            target_type = extract_torch_dtype(node)
+        assert isinstance(target_type, torch.dtype), type(target_type)
+
+        # define cast node
+        in_type: int = extract_circle_dtype(input)
+        out_type: int = to_circle_dtype(target_type)
+        inputs = [input]
+        outputs = [node]
+        operator = self.define_cast_node(inputs, outputs, in_type, out_type)
+
+        # TODO Support layout conversion
+
+        return operator

tico/serialize/operators/op_unsqueeze.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.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 UnSqueezeArgs
+
+
+@register_node_visitor
+class UnsqueezeVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.unsqueeze.default,
+        torch.ops.aten.unsqueeze_copy.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.EXPAND_DIMS, self._op_codes
+        )
+
+        args = UnSqueezeArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        dim = args.dim
+
+        dim_i32 = circle_legalize_dtype_to(dim, dtype=torch.int32)
+        inputs = [input, dim_i32]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.ExpandDimsOptions
+        )
+        option = circle.ExpandDimsOptions.ExpandDimsOptionsT()
+
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_view.py

@@ -0,0 +1,74 @@
+# 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.validate_args_kwargs import ViewArgs
+
+
+@register_node_visitor
+class ViewVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.view,
+        torch.ops.aten.view.default,
+        torch.ops.aten.view_copy.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 = ViewArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        size = args.size
+
+        assert is_const(size), type(size)
+
+        if isinstance(size, int):
+            raise Exception("scalar size conversion is not supported yet.")
+
+        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_where.py

@@ -0,0 +1,82 @@
+# 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 WhereSelfArgs
+
+
+@register_node_visitor
+class WhereVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.where.self]
+
+    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.SELECT_V2,
+            self._op_codes,
+        )
+
+        args = WhereSelfArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        condition = args.condition
+        input = args.input
+        other = args.other
+
+        result_true_dtype = (
+            extract_torch_dtype(input)
+            if isinstance(input, torch.fx.node.Node)
+            else input.dtype  # type: ignore[union-attr]
+        )
+        result_false_dtype = (
+            extract_torch_dtype(other)
+            if isinstance(other, torch.fx.node.Node)
+            else other.dtype  # type: ignore[union-attr]
+        )
+
+        if result_true_dtype != result_false_dtype:
+            raise RuntimeError(
+                f"Data type of arguments are not matched. result_true: {result_true_dtype}, result_false: {result_false_dtype}"
+            )
+
+        inputs = [condition, input, other]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = (
+            circle.BuiltinOptions.BuiltinOptions.SelectV2Options
+        )
+        option = circle.SelectV2Options.SelectV2OptionsT()
+
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/utils.py

@@ -0,0 +1,51 @@
+# 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
+
+from circle_schema import circle
+
+from tico.serialize.operators.hashable_opcode import OpCode
+
+
+def create_builtin_opcode(opcode: int) -> OpCode:
+    op_code = OpCode()
+    # deprecatedBuiltinCode is int8, so its maximum value is 127
+    # (127 is reserved as a placeholder for greater opcodes)
+    # opcode greater than 127 is saved in builtinCode
+    op_code.deprecatedBuiltinCode = min(127, opcode)
+    op_code.builtinCode = opcode
+    op_code.version = 1
+    return op_code
+
+
+def get_op_index(opcode: int, opcode_map: Dict[OpCode, int]) -> int:
+    op_code = create_builtin_opcode(opcode)
+    if op_code not in opcode_map:
+        op_index = len(opcode_map)
+        opcode_map[op_code] = op_index
+    else:
+        op_index = opcode_map[op_code]
+    return op_index
+
+
+# TODO Move this to CircleSubGraph
+def create_builtin_operator(
+    graph, op_index: int, inputs: List, outputs: List
+) -> circle.Operator.OperatorT:
+    operator = circle.Operator.OperatorT()
+    operator.opcodeIndex = op_index
+    operator.inputs = [graph.get_tid(input) for input in inputs]
+    operator.outputs = [graph.get_tid(output) for output in outputs]
+    return operator

tico/serialize/pack.py

@@ -0,0 +1,35 @@
+# 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 numpy as np
+
+
+def pack_buffer(flat_data: np.ndarray, dtype: str) -> np.ndarray:
+    assert len(flat_data.shape) == 1
+
+    if dtype == "uint4":
+        if flat_data.dtype != np.uint8:
+            raise RuntimeError("uint4 data should be saved in uint8.")
+
+        numel = flat_data.shape[0]
+        packed = np.zeros((numel + 1) // 2, dtype=np.uint8)
+        for i in range(numel):
+            assert flat_data[i] >= 0 and flat_data[i] <= 15
+            if i % 2 == 0:
+                packed[i // 2] = flat_data[i]
+            else:
+                packed[i // 2] |= flat_data[i] << 4
+        return packed
+    else:
+        raise NotImplementedError(f"NYI dtype: {dtype}")

tico/serialize/quant_param.py

@@ -0,0 +1,42 @@
+# 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.
+
+"""
+This is a key for torch.fx.Node's meta dict to save QuantParam
+
+QuantParam can be retrieved as node.meta[QPARAM_KEY]
+"""
+QPARAM_KEY = "_quantization_parameters_"
+
+from dataclasses import dataclass
+from typing import List, Optional
+
+import torch
+
+
+@dataclass
+class QuantParam:
+    scale: Optional[List[float]] = None
+    zero_point: Optional[List[int]] = None
+    quantized_dimension: Optional[int] = None
+    min: Optional[List[float]] = None
+    max: Optional[List[float]] = None
+    # NOTE We define dtype as a string to easily extend new dtypes (ex: uint4)
+    dtype: str = ""
+
+
+def to_qparam_dtype(dtype: torch.dtype) -> str:
+    str_type = str(dtype)
+    assert str_type.startswith("torch.")
+    return str_type[6:]

tico/utils/__init__.py

@@ -0,0 +1 @@
+# DO NOT REMOVE THIS FILE