tico 0.1.0__py3-none-any.whl

tico/serialize/circle_serializer.py

@@ -0,0 +1,240 @@
+# 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 operator
+from typing import Dict
+
+import flatbuffers
+import torch
+from circle_schema import circle
+from torch.export.exported_program import (
+    ConstantArgument,
+    ExportedProgram,
+    InputKind,
+    TensorArgument,
+)
+
+from tico.serialize.circle_mapping import to_circle_dtype
+from tico.serialize.operators import *
+from tico.serialize.circle_graph import CircleModel, CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import get_node_visitors
+from tico.utils import logging
+from tico.utils.serialize import finalise_tensor_names
+
+
+multiple_output_ops = [
+    torch.ops.aten.split_with_sizes.default,
+    torch.ops.aten.max.dim,
+]
+
+# Build circle model from ExportedProgram
+# Return raw bytes of circle model
+def build_circle(edge_program: ExportedProgram) -> bytes:
+    logger = logging.getLogger(__name__)
+
+    builder = flatbuffers.Builder()
+
+    # Init Model
+    model = CircleModel()
+
+    # Add empty buffer at the front (convention)
+    model.add_buffer(circle.Buffer.BufferT())
+
+    # Create an empty subgraph (assume a single subgraph)
+    graph = CircleSubgraph(model)
+
+    # Export tensors
+    logger.debug("---------------Export tensors--------------")
+    buf_name_to_data = {name: buf for name, buf in edge_program.named_buffers()}
+    for node in edge_program.graph.nodes:
+        if node.op == "call_function":
+            if node.target in multiple_output_ops:
+                continue
+            node_val = node.meta["val"]
+            if node_val.layout != torch.strided:
+                raise RuntimeError(
+                    f"Only support dense tensors (node layout: {node_val.layout})"
+                )
+            graph.add_tensor_from_node(node)
+            logger.debug(f"call_function: {node.name} tensor exported.")
+
+        # placeholder: function input (including parameters, buffers, constant tensors)
+        elif node.op == "placeholder":
+            # placeholder invariants
+            assert node.args is None or len(node.args) == 0  # Not support default param
+
+            # parameters
+            if node.name in edge_program.graph_signature.inputs_to_parameters:
+                param_name = edge_program.graph_signature.inputs_to_parameters[
+                    node.name
+                ]
+                param_data = edge_program.state_dict[param_name]
+
+                assert isinstance(
+                    param_data, torch.Tensor
+                ), "Expect parameters to be a tensor"
+                param_value = param_data.cpu().detach().numpy()
+
+                graph.add_tensor_from_node(node, param_value)
+                logger.debug(f"placeholder(param): {node.name} tensor exported.")
+            elif node.name in edge_program.graph_signature.inputs_to_buffers:
+                buffer_name = edge_program.graph_signature.inputs_to_buffers[node.name]
+                assert buffer_name in buf_name_to_data
+                buffer_data = buf_name_to_data[buffer_name]
+                assert isinstance(
+                    buffer_data, torch.Tensor
+                ), "Expect buffers to be a tensor"
+                buffer_value = buffer_data.cpu().detach().numpy()
+
+                graph.add_tensor_from_node(node, buffer_value)
+                logger.debug(f"placeholder(buffer): {node.name} tensor exported.")
+            elif (
+                node.name
+                in edge_program.graph_signature.inputs_to_lifted_tensor_constants
+            ):
+                ctensor_name = (
+                    edge_program.graph_signature.inputs_to_lifted_tensor_constants[
+                        node.name
+                    ]
+                )
+                ctensor_data = edge_program.constants[ctensor_name]
+
+                assert isinstance(
+                    ctensor_data, torch.Tensor
+                ), "Expect constant tensor to be a tensor"
+                ctensor_value = ctensor_data.cpu().detach().numpy()
+
+                graph.add_tensor_from_node(node, ctensor_value)
+                logger.debug(
+                    f"placeholder(constant tensor): {node.name} tensor exported."
+                )
+            else:
+                user_inputs = [
+                    specs
+                    for specs in edge_program.graph_signature.input_specs
+                    if specs.kind == InputKind.USER_INPUT
+                ]
+                constant_inputs = [
+                    specs
+                    for specs in user_inputs
+                    if isinstance(specs.arg, ConstantArgument)
+                ]
+                name_to_value = {
+                    specs.arg.name: specs.arg.value for specs in constant_inputs
+                }
+                # NoneType ConstantArgument is ignored.
+                if node.name in name_to_value and name_to_value[node.name] == None:
+                    continue
+                graph.add_tensor_from_node(node)
+                logger.debug(f"placeholder: {node.name} tensor exported.")
+
+        # get_attr: retrieve parameter
+        elif node.op == "get_attr":
+            # node.name: Place where fetched attribute is saved
+            # node.target: Attribute in the module
+            attr_tensor = getattr(node.graph.owning_module, node.target)
+            assert isinstance(attr_tensor, torch.Tensor)
+
+            graph.add_tensor_from_scratch(
+                prefix=node.name,
+                shape=list(attr_tensor.shape),
+                dtype=to_circle_dtype(attr_tensor.dtype),
+                source_node=node,
+            )
+
+            logger.debug(f"get_attr: {node.name} tensor exported.")
+
+        # output: function output
+        elif node.op == "output":
+            # output node itself does not need a buffer
+            # argument of output node is assumed to be exported beforehand
+            for output in node.args[0]:
+                if isinstance(output, torch.fx.Node):
+                    assert graph.has_tensor(output.name)
+            continue
+
+        # call_method: call method
+        elif node.op == "call_method":
+            raise AssertionError("Not yet implemented")
+
+        # call_module: call 'forward' of module
+        elif node.op == "call_module":
+            raise AssertionError("Not yet implemented")
+
+        else:
+            # Add more if fx.Node is extended
+            raise AssertionError(f"Unknown fx.Node op {node.op}")
+
+    # Register inputs
+    logger.debug("---------------Register inputs--------------")
+    for in_spec in edge_program.graph_signature.input_specs:
+        if in_spec.kind != InputKind.USER_INPUT:
+            continue
+        # NoneType ConstantArgument is ignored.
+        if isinstance(in_spec.arg, ConstantArgument) and in_spec.arg.value == None:
+            continue
+        arg_name = in_spec.arg.name
+        graph.add_input(arg_name)
+        logger.debug(f"Registered input: {arg_name}")
+
+    # Register outputs
+    logger.debug("---------------Register outputs--------------")
+    for user_output in edge_program.graph_signature.user_outputs:
+        if user_output == None:
+            logger.debug(f"Ignore 'None' output")
+            continue
+
+        graph.add_output(user_output)
+        logger.debug(f"Registered output: {user_output}")
+
+    # Export operators
+    logger.debug("---------------Export operators--------------")
+    op_codes: Dict[OpCode, int] = {}
+    visitors = get_node_visitors(op_codes, graph)
+    for node in edge_program.graph.nodes:
+        if node.op != "call_function":
+            continue
+
+        opcode = node.target
+        if opcode == operator.getitem:
+            continue
+        if opcode not in visitors:
+            raise RuntimeError(f"{opcode} is not yet supported")
+        circle_op = visitors[opcode].define_node(node)
+
+        if circle_op:
+            graph.add_operator(circle_op)
+            logger.debug(f"call_function: {node.name} ({opcode}) Op exported.")
+
+    # Register subgraph
+    finalise_tensor_names(graph)
+    model.subgraphs.append(graph)
+
+    # Encode operator codes
+    model.operatorCodes = [
+        code for code, _ in sorted(op_codes.items(), key=lambda x: x[1])
+    ]
+
+    # Description
+    model.description = "circle"
+
+    # Set version
+    model.version = 0
+
+    # Finish model
+    builder.Finish(model.Pack(builder), "CIR0".encode("utf8"))
+    buf = builder.Output()
+
+    return bytes(buf)

tico/serialize/operators/__init__.py

@@ -0,0 +1,28 @@
+# 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 glob
+from os.path import basename, dirname, isfile, join
+
+from tico.utils.register_custom_op import RegisterOps
+
+
+# Register custom ops to torch namespace
+RegisterOps()
+
+# Load all modules in the current directory
+modules = glob.glob(join(dirname(__file__), "*.py"))
+__all__ = [
+    basename(f)[:-3] for f in modules if isfile(f) and not f.endswith("__init__.py")
+]

tico/serialize/operators/hashable_opcode.py

@@ -0,0 +1,43 @@
+# 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 circle_schema import circle
+
+
+class OpCode(circle.OperatorCode.OperatorCodeT):
+    """
+    Wrapper class for operator code in circle schema
+    This implements __eq__ and __hash__ for use with dict()
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def __eq__(self, other):
+        if self.version != other.version:
+            return False
+
+        if self.builtinCode == circle.BuiltinOperator.BuiltinOperator.CUSTOM:
+            return self.customCode == other.customCode
+
+        return self.builtinCode == other.builtinCode
+
+    def __hash__(self):
+        val = (
+            self.deprecatedBuiltinCode,
+            self.customCode,
+            self.version,
+            self.builtinCode,
+        )
+        return hash(val)

tico/serialize/operators/node_visitor.py

@@ -0,0 +1,80 @@
+# 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, Type, 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.operators.hashable_opcode import OpCode
+
+
+class NodeVisitor:
+    """
+    Node visitor for lowering edge IR to circle
+    """
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        # For setting opcode index in circle model
+        # This is updated during serialization
+        self._op_codes = op_codes
+        self.graph = graph
+
+    # Define circle model operator
+    def define_node(
+        self,
+        node: torch.fx.node.Node,
+    ) -> circle.Operator.OperatorT:
+        raise NotImplementedError("NodeVisitor must be extended.")
+
+
+# container for all node visitors
+_node_visitor_dict: Dict[torch._ops.OpOverload, Type[NodeVisitor]] = {}
+
+
+# Decorator for each visitor
+def register_node_visitor(visitor):
+    for target in visitor.target:
+        _node_visitor_dict[target] = visitor
+    return visitor
+
+
+def get_node_visitor(target: torch._ops.OpOverload) -> Type[NodeVisitor]:
+    """
+    Get a single node visitor (for unittest purpose)
+    """
+    _visitor = _node_visitor_dict.get(target, None)
+
+    if not _visitor:
+        raise LookupError(f"NodeVisitor for {target} is not registered")
+
+    return _visitor
+
+
+# Get all node visitors
+def get_node_visitors(
+    op_codes: Dict[OpCode, int], graph: CircleSubgraph
+) -> Dict[torch._ops.OpOverload, NodeVisitor]:
+    node_visitors = {}
+    for target, visitor in _node_visitor_dict.items():
+        node_visitors[target] = visitor(op_codes, graph)
+
+    return node_visitors
+
+
+def get_support_targets():
+    return _node_visitor_dict.keys()

tico/serialize/operators/op_abs.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 AbsArgs
+
+
+@register_node_visitor
+class AbsVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [torch.ops.aten.abs.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.ABS, self._op_codes
+        )
+
+        args = AbsArgs(*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_add.py

@@ -0,0 +1,69 @@
+# 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 AddTensorArgs
+
+
+@register_node_visitor
+class AddVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.add.Tensor,
+        torch.ops.aten.add.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 = AddTensorArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+        input = args.input
+        other = args.other
+
+        inputs = [input, other]
+        outputs = [node]
+
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.ADD, self._op_codes
+        )
+
+        inputs = [input, other]
+        outputs = [node]
+
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+
+        # Op-specific option
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.AddOptions
+        option = circle.AddOptions.AddOptionsT()
+        option.fusedActivationFunction = (
+            circle.ActivationFunctionType.ActivationFunctionType.NONE
+        )
+        option.potScaleInt16 = False
+        operator.builtinOptions = option
+
+        return operator

tico/serialize/operators/op_alias_copy.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.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 AliasCopyArgs
+
+
+@register_node_visitor
+class AliasCopyVisitor(NodeVisitor):
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.aten.alias.default,
+        torch.ops.aten.alias_copy.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 = AliasCopyArgs(*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