tico 0.1.0.dev250411__py3-none-any.whl

tico/passes/remove_redundant_to_copy.py

@@ -0,0 +1,84 @@
+# 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 Union
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.passes import ops
+from tico.serialize.circle_mapping import extract_torch_dtype
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.validate_args_kwargs import ToCopyArgs, ToDtypeArgs, ToDtypeLayoutArgs
+
+
+@trace_graph_diff_on_pass
+class RemoveRedundantToCopy(PassBase):
+    """
+    This pass removes redundant `aten._to_copy` operators.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        modified = False
+        for node in graph.nodes:
+            if not node.op == "call_function":
+                continue
+
+            if not node.target in ops.aten.to_copy:
+                continue
+
+            args: Union[ToCopyArgs, ToDtypeArgs, ToDtypeLayoutArgs]
+            if node.target == torch.ops.aten._to_copy.default:
+                args = ToCopyArgs(*node.args, **node.kwargs)
+            elif node.target == torch.ops.aten.to.dtype:
+                args = ToDtypeArgs(*node.args, **node.kwargs)
+            elif node.target == torch.ops.aten.to.dtype_layout:
+                args = ToDtypeLayoutArgs(*node.args, **node.kwargs)
+            else:
+                raise NotImplementedError(
+                    f"Unsupported to_copy operator: {node.target}"
+                )
+
+            input_ = args.input
+            # 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 hasattr(args, "layout") and args.layout is not None:
+                if args.layout != input_.meta["val"].layout:
+                    continue
+
+            if hasattr(args, "dtype") and args.dtype is not None:
+                target_dtype = args.dtype
+                input_dtype = extract_torch_dtype(input_)
+                if input_dtype != target_dtype:
+                    continue
+
+            node.replace_all_uses_with(input_, propagate_meta=False)
+
+            modified = True
+            logger.debug(f"{node.name} is replaced with {input_.name}")
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)

tico/passes/restore_linear.py

@@ -0,0 +1,113 @@
+# 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 torch
+from torch.export import ExportedProgram
+
+from tico.utils import logging
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+
+
+@trace_graph_diff_on_pass
+class RestoreLinear(PassBase):
+    """
+    Linear Op is decomposed to multiple Ops in core aten
+    This pass restores linear Ops. For example,
+
+    Before)
+
+        bias   input   weight         input    weight
+        |        |       |              |        |
+        |        |   permute_copy       |    permute_copy
+        |        V       |              |        |
+        +----> addmm <---+              |        V
+                                        +------> mm
+
+    After)
+
+        input  weight   bias          input    weight
+        |        |       |              |        |
+        |        |       |              |        |
+        |        V       |              |        V
+        +---> linear <---+              +----> linear
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        modified = False
+        for node in graph.nodes:
+            if not node.op == "call_function":
+                continue
+
+            if node.target == torch.ops.aten.addmm.default:
+                assert len(node.args) == 3
+                bias, input, permute = node.args
+                if permute.target not in [
+                    torch.ops.aten.permute.default,
+                    torch.ops.aten.t.default,
+                ]:
+                    continue
+
+                if permute.target == torch.ops.aten.permute_copy.default:
+                    dims = permute.args[1]
+                    if dims != [1, 0]:
+                        continue
+                weight = permute.args[0]
+
+                addmm_args = (input, weight, bias)
+                with graph.inserting_after(node):
+                    linear_node = graph.call_function(
+                        torch.ops.aten.linear.default, args=addmm_args
+                    )
+                    node.replace_all_uses_with(linear_node, propagate_meta=True)
+                    graph.erase_node(node)
+
+            elif node.target == torch.ops.aten.mm.default:
+                assert len(node.args) == 2
+                input, permute = node.args
+                if permute.target not in [
+                    torch.ops.aten.permute.default,
+                    torch.ops.aten.t.default,
+                ]:
+                    continue
+
+                if permute.target == torch.ops.aten.permute_copy.default:
+                    dims = permute.args[1]
+                    if dims != [1, 0]:
+                        continue
+                weight = permute.args[0]
+
+                mm_args = (input, weight)
+                with graph.inserting_after(node):
+                    linear_node = graph.call_function(
+                        torch.ops.aten.linear.default, args=mm_args
+                    )
+                    node.replace_all_uses_with(linear_node, propagate_meta=True)
+
+            else:
+                continue
+
+            modified = True
+            logger.debug(f"{node.name} is replaced with linear")
+
+        graph.eliminate_dead_code()
+        graph_module.recompile()
+        return PassResult(modified)

tico/passes/segment_index_select.py

@@ -0,0 +1,143 @@
+# 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 TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+
+from torch._export.utils import (
+    get_buffer,
+    get_lifted_tensor_constant,
+    get_param,
+    is_buffer,
+    is_lifted_tensor_constant,
+    is_param,
+)
+from torch.export import ExportedProgram
+
+from tico.passes import ops
+from tico.serialize.circle_graph import extract_shape
+from tico.utils import logging
+from tico.utils.graph import add_placeholder, is_single_value_tensor
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_const_diff_on_pass
+from tico.utils.validate_args_kwargs import IndexSelectArgs
+
+
+@trace_const_diff_on_pass
+class SegmentIndexSelectConst(PassBase):
+    """
+    Let's segment index_select with multiple const indices to index_select operators with one index.
+    WHY?
+      Gather(index, index_select, select, embedding, ...) operation with const indices can be lowered to slice by LowerToSlice pass.
+      For that, we need to split 'a index_select operator with multiple indice' to 'multiple index_select operators with one index'.
+      Note that NPU is not fully compatible with gather operation.
+
+    [before]
+            input
+                |
+            index_select.default, len(index) > 1
+                |
+            output
+
+    [after]
+
+            input
+                |
+                -------------------------------------------------
+                |                                                |
+            index_select.default, len(index) == 1  , ... , index_select.default, len(index) == 1
+                |                                                |
+                -------------------------------------------------
+                |
+            torch.concat (input=[index_select0, index_select1, ...], axis = dim)
+                |
+            output
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        modified = False
+        for node in graph.nodes:
+            if not node.op == "call_function":
+                continue
+
+            if not node.target in ops.aten.index_select:
+                continue
+
+            args = IndexSelectArgs(*node.args, **node.kwargs)
+            input = args.input
+            dim = args.dim
+            index = args.index
+
+            if isinstance(index, torch.fx.Node):
+                if is_lifted_tensor_constant(exported_program, index):
+                    index = get_lifted_tensor_constant(exported_program, index)  # type: ignore[assignment]
+                elif is_param(exported_program, index):
+                    index = get_param(exported_program, index)  # type: ignore[assignment]
+                elif is_buffer(exported_program, index):
+                    index = get_buffer(exported_program, index)  # type: ignore[assignment]
+                else:
+                    continue
+
+            if not isinstance(index, torch.Tensor):
+                continue
+
+            if is_single_value_tensor(index):
+                continue
+
+            if len(index) < 2:
+                continue
+
+            input_shape = extract_shape(input)
+            if dim < 0:
+                dim = dim % len(input_shape)
+
+            index_select_node_list = []
+            for i in index:
+                index_node = add_placeholder(
+                    exported_program, torch.tensor([i]), prefix="segm_index"
+                )
+                with graph.inserting_before(node):
+                    index_select_node = graph.call_function(
+                        torch.ops.aten.index_select.default,
+                        args=(input, dim, index_node),
+                    )
+                    index_select_node_list.append(index_select_node)
+
+            with graph.inserting_before(node):
+                concat_node = graph.call_function(
+                    torch.ops.aten.cat.default, args=(index_select_node_list, dim)
+                )
+
+                node.replace_all_uses_with(concat_node, propagate_meta=False)
+
+            modified = True
+            logger.debug(
+                f"{node.name} is replaced with {concat_node.name} and {index_select_node_list}"
+            )
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)

tico/pt2_to_circle.py

@@ -0,0 +1,101 @@
+# 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 argparse
+import os
+
+import torch
+import yaml
+
+from tico.config import CompileConfigBase, get_default_config
+
+from tico.utils.convert import convert_exported_module_to_circle
+
+
+def convert(
+    input: str,
+    output: str,
+    verbose: bool = False,
+    config: CompileConfigBase = get_default_config(),
+):
+    # TODO Check input and output
+
+    if verbose:
+        os.environ["TICO_LOG"] = "4"
+
+    exported_program = torch.export.load(input)
+    circle_program = convert_exported_module_to_circle(exported_program, config=config)
+    circle_binary = circle_program
+    with open(output, "wb") as f:
+        f.write(circle_binary)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-i",
+        "--input",
+        required=True,
+        help="provide a path to .pt2 model.",
+    )
+
+    parser.add_argument(
+        "-o",
+        "--output",
+        required=True,
+        help="provide a path to .circle model.",
+    )
+
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="print logs.",
+    )
+
+    parser.add_argument(
+        "-c",
+        "--config",
+        required=False,
+        help="provide a path to config file.",
+    )
+
+    args = parser.parse_args()
+
+    if args.config:
+        with open(args.config) as f:
+            config_dict = yaml.safe_load(f)
+
+            version = config_dict.get("version", None)
+            latest_version = "1.0"
+
+            if version is None:
+                raise ValueError(
+                    f"'version' field must be provided in the config file. (lastest: {latest_version})"
+                )
+
+            if version == "1.0":
+                from tico.config.v1 import CompileConfigV1
+
+                config = CompileConfigV1.from_dict(config_dict)
+            else:
+                raise ValueError(
+                    f"Unsupported version '{version}'. (lastest: {latest_version})"
+                )
+
+    convert(args.input, args.output, args.verbose, config)
+
+
+if __name__ == "__main__":
+    main()

tico/serialize/__init__.py

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