tico 0.1.0.dev250609__py3-none-any.whl → 0.1.0.dev250610__py3-none-any.whl

tico/__init__.py

@@ -21,7 +21,7 @@ from tico.config import CompileConfigV1, get_default_config
 from tico.utils.convert import convert, convert_from_exported_program, convert_from_pt2
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250609"
+__version__ = "0.1.0.dev250610"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/passes/fuse_leading_unsqueeze_reshape.py

@@ -0,0 +1,110 @@
+# 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 Sequence
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.passes import ops
+from tico.serialize.circle_mapping import extract_shape
+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.utils import is_single_use_target_node
+from tico.utils.validate_args_kwargs import PermuteArgs, ReshapeArgs
+
+
+def _is_leading_unsqueeze(target: Sequence[int], permuted: Sequence[int]) -> bool:
+    """
+    True if `target` == [1]*k + permuted,  k>=1.
+    """
+    k = len(target) - len(permuted)
+    return (
+        k > 0 and all(d == 1 for d in target[:k]) and list(target[k:]) == list(permuted)
+    )
+
+
+@trace_graph_diff_on_pass
+class FuseLeadingUnsqueezeReshape(PassBase):
+    """
+    Fuse reshape → permute → reshape where the second reshape only
+    prepends one-sized dims (unsqueeze) to the permuted tensor.
+
+    [BEFORE]
+        x - aten.reshape(s1) - aten.permute(p) - aten.reshape([1]*k + p(s1))
+    [AFTER]
+        x - aten.reshape([1]*k + s1) - aten.permute(list(range(k)) + [d+k for d in p])
+    """
+
+    def call(self, ep: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        gm = ep.graph_module
+        graph = gm.graph
+        modified = False
+        for reshape_back in graph.nodes:
+            if (
+                reshape_back.op != "call_function"
+                or reshape_back.target not in ops.aten.reshape
+            ):
+                continue
+            reshape_back_args = ReshapeArgs(*reshape_back.args, **reshape_back.kwargs)  # type: ignore[arg-type]
+            permute, reshape_back_size = reshape_back_args.input, reshape_back_args.size
+
+            if not is_single_use_target_node(permute, ops.aten.permute):
+                continue
+            permute_args = PermuteArgs(*permute.args, **permute.kwargs)  # type: ignore[arg-type]
+            reshape_front, permute_dims = permute_args.input, permute_args.dims
+
+            if not is_single_use_target_node(reshape_front, ops.aten.reshape):
+                continue
+            reshape_front_args = ReshapeArgs(*reshape_front.args, **reshape_front.kwargs)  # type: ignore[arg-type]
+            reshape_front_input, reshape_front_size = (
+                reshape_front_args.input,
+                reshape_front_args.size,
+            )
+
+            # ---- condition: only leading unsqueeze ------------------------
+            back_shape = extract_shape(reshape_back)
+            permute_shape = extract_shape(permute)
+
+            if not _is_leading_unsqueeze(back_shape, permute_shape):
+                continue
+
+            # ---- create new reshape & new permute -------------------------
+            k = len(back_shape) - len(permute_shape)
+            with graph.inserting_before(permute):
+                new_shape = [1] * k + list(reshape_front_size)
+                r_new = graph.call_function(
+                    torch.ops.aten.reshape.default,
+                    args=(reshape_front_input, new_shape),
+                )
+                new_p_dims = list(range(k)) + [
+                    d + k for d in permute_dims
+                ]  # shift by k
+                p_new = graph.call_function(
+                    torch.ops.aten.permute.default, args=(r_new, new_p_dims)
+                )
+
+            reshape_back.replace_all_uses_with(p_new, propagate_meta=True)
+            modified = True
+            logger.debug(f"{reshape_back.name} is fused to {r_new.name}")
+
+        if modified:
+            graph.eliminate_dead_code()
+            graph.lint()
+            gm.recompile()
+
+        return PassResult(modified)

tico/utils/convert.py

@@ -51,6 +51,7 @@ from tico.passes.decompose_grouped_conv2d import DecomposeGroupedConv2d
 from tico.passes.decompose_slice_scatter import DecomposeSliceScatter
 from tico.passes.extract_dtype_kwargs import ExtractDtypeKwargsPass
 from tico.passes.fill_meta_val import FillMetaVal
+from tico.passes.fuse_leading_unsqueeze_reshape import FuseLeadingUnsqueezeReshape
 from tico.passes.fuse_redundant_reshape_to_mean import FuseRedundantReshapeToMean
 from tico.passes.legalize_causal_mask_value import LegalizeCausalMaskValue
 from tico.passes.legalize_predefined_layout_operators import (
@@ -225,6 +226,7 @@ def convert_exported_module_to_circle(
             LowerPow2ToMul(),
             ConvertConv1dToConv2d(),
             *LowerToSlicePasses(),
+            FuseLeadingUnsqueezeReshape(),
         ]
     )
     circle_legalize.run(exported_program)

tico/utils/utils.py

@@ -378,3 +378,31 @@ def broadcastable(
         if dim_a != 1 and dim_b != 1 and dim_a != dim_b:
             return False
     return True
+
+
+def is_single_use_target_node(
+    node: torch.fx.Node, target_ops: list[torch._ops.OpOverload] | torch._ops.OpOverload
+):
+    """
+    Check whether a given node is a `call_function` node that matches one of the specified targets
+    and is used by only one other node.
+
+    Args:
+        node (torch.fx.Node): The node to check.
+        target_ops (Iterable[Callable]): A list or set of target operations to match (e.g., ops.aten.reshape).
+
+    Returns:
+        bool: True if the node is a call_function, its target is in `target_ops`, and it has exactly one user.
+    """
+    if not isinstance(target_ops, list):
+        target_ops = [target_ops]
+    assert all(isinstance(t, torch._ops.OpOverload) for t in target_ops), target_ops
+
+    if node.op != "call_function":
+        return False
+    if node.target not in target_ops:
+        return False
+    if len(node.users) != 1:
+        return False
+
+    return True

{tico-0.1.0.dev250609.dist-info → tico-0.1.0.dev250610.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tico
-Version: 0.1.0.dev250609
+Version: 0.1.0.dev250610
 Summary: Convert exported Torch module to circle
 Home-page: UNKNOWN
 License: UNKNOWN

{tico-0.1.0.dev250609.dist-info → tico-0.1.0.dev250610.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=mEQHwb4j_KQrsDMO6gp2USb5RCYNhd0MOxL4P_EGMW8,1743
+tico/__init__.py,sha256=PHlL6DvW31YH0rgvY-y2LlEJDbNJSK7FO1KX1O5BGMM,1743
 tico/pt2_to_circle.py,sha256=PPmFNw20jw2Z2VyM3ln9pX__jTzBOAZiv0gT5a-p-Y8,2666
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=anwOiJFkUxUi7Cef573JgQcjk6S-FSi6O_TLjYASW-g,1244
@@ -75,6 +75,7 @@ tico/passes/decompose_grouped_conv2d.py,sha256=KJhH6PX7l9k9T8KBV8JDAvaSfJuUnRo_j
 tico/passes/decompose_slice_scatter.py,sha256=ko9p8v-zY5rOx4aSpWomwSdSWb1lIF32gnU7ik5xgII,5604
 tico/passes/extract_dtype_kwargs.py,sha256=hfGJ_GfZULbBmLif2AJkhPHVifhucxBiLoQI862Yejk,4303
 tico/passes/fill_meta_val.py,sha256=Xbam6Aq90ZfWItZw1dgLIwH_q8RCiU5JodKNqkj-ink,1797
+tico/passes/fuse_leading_unsqueeze_reshape.py,sha256=1hqX9urmV4kQXGWCEJ22XNH_3M2Og5v09N2wwV87FoQ,4304
 tico/passes/fuse_redundant_reshape_to_mean.py,sha256=SzHLLL2Yfj4k1c2L5i4PVI9EFUilHRfIS6S-hahKFCM,3702
 tico/passes/legalize_causal_mask_value.py,sha256=KZc_UPk7CGPXO35JOu6dVrOzRJx-ZpJoVvuzz-GvQek,4080
 tico/passes/legalize_predefined_layout_operators.py,sha256=N2TtJInjSTk-E5afnkDXXbo9v4zTM7yzsjna3VoihMw,15895
@@ -177,7 +178,7 @@ tico/serialize/operators/op_view.py,sha256=5EMww-ve17Vm9XPuV03Tn7vJsjpU2J8U4d_FO
 tico/serialize/operators/op_where.py,sha256=qZDFKVQ2u8HB0Jjpto_1A4g-jfpEprAbBT3PmIGzdoY,2855
 tico/serialize/operators/utils.py,sha256=lXGpEJW1h8U_-gfc6EWjvvSiq3yJ9P-v1v3EMRT_pSk,2954
 tico/utils/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
-tico/utils/convert.py,sha256=KCllPnvQ8bjEYR1yI72s9aNBp7Py1CzIEEpYSYZcu60,11684
+tico/utils/convert.py,sha256=nbpBJue9_ezb8Dj9L9dWcNeVlZd0soGJu_fvrL6-c5c,11810
 tico/utils/define.py,sha256=Ypgp7YffM4pgPl4Zh6TmogSn1OxGBMRw_e09qYGflZk,1467
 tico/utils/diff_graph.py,sha256=_eDGGPDPYQD4b--MXX0DLoVgSt_wLfNPt47UlolLLR4,5272
 tico/utils/errors.py,sha256=f3csJjgbXG9W1aHhqEcou008Aor19W57X8oT5Hx8w1M,954
@@ -189,15 +190,15 @@ tico/utils/passes.py,sha256=kGmDe__5cPaO6i5EDAoXSVe6yXEoX9hAny4ROb3ZEmQ,2409
 tico/utils/register_custom_op.py,sha256=iRQvdqlBqrJxq_pNkvJyDIJD_SYtCUl88wwbbuvSwlk,22952
 tico/utils/serialize.py,sha256=AQXMBOLu-Kg2Rn-qbqsAtHndjZAZIavlKA0QFgJREHM,1420
 tico/utils/trace_decorators.py,sha256=ddLIiKQfSaQrxgF1kNpwjFTQnXENzeSfcr1kuAW4jGI,3221
-tico/utils/utils.py,sha256=7t4U3Bs7SV7y3PCdgE9bEAAAI2u1mzQEUP81oP_XHHU,12334
+tico/utils/utils.py,sha256=A5YeUXIjSY5HnlmtnMfWm6wx3i6Feq5dSra1NJsZV2o,13253
 tico/utils/validate_args_kwargs.py,sha256=P4aMnr9EhNCtc_AgJPpuezfQbqFfDn0lhJSWqmumLZ8,25054
 tico/utils/mx/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/utils/mx/elemwise_ops.py,sha256=V6glyAHsVR1joqpsgnNytatCD_ew92xNWZ19UFDoMTA,10281
 tico/utils/mx/formats.py,sha256=uzNWyu-1onUlwQfX5cZ6fZSUfHMRqorper7_T1k3jfk,3404
 tico/utils/mx/mx_ops.py,sha256=RcfUTYVi-wilGB2sC35OeARdwDqnixv7dG5iyZ-fQT8,8555
-tico-0.1.0.dev250609.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250609.dist-info/METADATA,sha256=ZpNUuShOOV1hMrNe84u05IoeX4UZn1fj1_TEU4ZlI1Y,8633
-tico-0.1.0.dev250609.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250609.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250609.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250609.dist-info/RECORD,,
+tico-0.1.0.dev250610.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250610.dist-info/METADATA,sha256=9nX12pxSQ-4h2gPPm0vYX0EW3wkn00HtmLF1E-xIabE,8633
+tico-0.1.0.dev250610.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250610.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250610.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250610.dist-info/RECORD,,