tico 0.1.0.dev250608__py3-none-any.whl → 0.1.0.dev250609__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.dev250608"
+__version__ = "0.1.0.dev250609"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/passes/ops.py

@@ -54,7 +54,10 @@ class AtenOps:
         self.reshape = [torch.ops.aten.reshape.default]
         self.select = [torch.ops.aten.select_copy.int, torch.ops.aten.select.int]
         self.slice = [torch.ops.aten.slice.Tensor, torch.ops.aten.slice_copy.Tensor]
-        self.softmax = [torch.ops.aten._softmax.default]
+        self.softmax = [
+            torch.ops.aten._softmax.default,
+            torch.ops.aten._safe_softmax.default,
+        ]
         self.squeeze = [torch.ops.aten.squeeze.dims, torch.ops.aten.squeeze_copy.dims]
         self.to_copy = [
             torch.ops.aten._to_copy.default,

tico/passes/remove_redundant_permute.py

@@ -24,6 +24,22 @@ from tico.serialize.circle_mapping import extract_shape, extract_stride
 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 PermuteArgs
+
+
+def _compose_permutation(dims1: list[int], dims2: list[int]):
+    """
+    Compose two permutation vectors.
+
+    Given y = x.permute(dims1) and z = y.permute(dims2),
+    the overall permutation p = dims2 ∘ dims1 is
+
+        p[i] = dims1[dims2[i]]
+    """
+    assert len(dims1) == len(
+        dims2
+    ), f"len(dims1): {len(dims1)}, len(dims2): {len(dims2)}"
+    return [dims1[i] for i in dims2]
 
 
 def passes():
@@ -45,9 +61,13 @@ class RemoveRedundantPermutePattern1(PassBase):
     def call(self, exported_program: ExportedProgram) -> PassResult:
         """
         [BEFORE]
-            (AxBxC) - aten.permute - aten.permute - (AxBxC)
+            (AxBxC) - aten.permute_1 - aten.permute_2 - (OUT_SHAPE)
         [AFTER]
-            (AxBxC)
+            if OUT_SHAPE == (AxBxC):
+                (AxBxC)
+            else:
+                (AxBxC) - aten.permute (fused dims) - (OUT_SHAPE)
+
         """
         logger = logging.getLogger(__name__)
 
@@ -61,39 +81,36 @@ class RemoveRedundantPermutePattern1(PassBase):
                 continue
             if len(permute2.users) != 1:
                 continue
-            assert len(permute2.args) == 2
-            permute1, permute2_dims = permute2.args
-            assert isinstance(permute1, torch.fx.Node), type(permute1)
-            assert isinstance(permute2_dims, list), type(permute2_dims)
-            for dim in permute2_dims:
-                assert isinstance(dim, int), type(dim)
+            permute2_args = PermuteArgs(*permute2.args, **permute2.kwargs)  # type: ignore[arg-type]
+            permute1, permute2_dims = permute2_args.input, permute2_args.dims
 
             if not permute1.target in ops.aten.permute:
                 continue
             if len(permute1.users) != 1:
                 continue
-            assert len(permute1.args) == 2
-            permute1_input, permute1_dims = permute1.args
-            assert isinstance(permute1_input, torch.fx.Node), type(permute1_input)
-            assert isinstance(permute1_dims, list), type(permute1_dims)
-            for dim in permute1_dims:
-                assert isinstance(dim, int), type(dim)
-
-            # shape
-            permute1_input_shape = extract_shape(permute1_input)
-            permute2_shape = extract_shape(permute2)
-            if permute1_input_shape != permute2_shape:
-                continue
-            # stride
-            permute1_input_stride = extract_stride(permute1_input)
-            permute2_stride = extract_stride(permute2)
-            if permute1_input_stride != permute2_stride:
-                continue
-
-            permute2.replace_all_uses_with(permute1_input, propagate_meta=False)
-
+            permute1_args = PermuteArgs(*permute1.args, **permute1.kwargs)  # type: ignore[arg-type]
+            permute1_input, permute1_dims = permute1_args.input, permute1_args.dims
+
+            fused_dims = _compose_permutation(permute1_dims, permute2_dims)
+            identity = list(range(len(fused_dims)))
+
+            if fused_dims == identity:
+                # shape
+                permute1_input_shape = extract_shape(permute1_input)
+                permute2_shape = extract_shape(permute2)
+                assert permute1_input_shape == permute2_shape
+
+                permute2.replace_all_uses_with(permute1_input, propagate_meta=False)
+                logger.debug(f"{permute1.name} and {permute2.name} are removed.")
+            else:
+                with graph.inserting_after(permute2):
+                    new_args = (permute1_input, fused_dims)
+                    fused_permute = graph.call_function(
+                        torch.ops.aten.permute.default, args=new_args
+                    )
+                    permute2.replace_all_uses_with(fused_permute, propagate_meta=True)
+                    logger.debug(f"{permute1.name} and {permute2.name} are fused.")
             modified = True
-            logger.debug(f"{permute1.name} and {permute2.name} are removed.")
 
         graph.eliminate_dead_code()
         graph.lint()

tico/passes/remove_redundant_reshape.py

@@ -24,11 +24,12 @@ 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 set_new_meta_val
+from tico.utils.utils import broadcastable, set_new_meta_val
 from tico.utils.validate_args_kwargs import (
     AddTensorArgs,
     PermuteArgs,
     ReshapeArgs,
+    SafeSoftmaxArgs,
     SoftmaxArgs,
 )
 
@@ -253,15 +254,14 @@ class RemoveRedundantReshapePattern3(PassBase):
                 continue
             if not softmax.target in ops.aten.softmax:
                 continue
-            softmax_args = SoftmaxArgs(*softmax.args, **softmax.kwargs)  # type: ignore[arg-type]
-            add, softmax_dim, softmax_half_to_float = (
+            if softmax.target == torch.ops.aten._softmax.default:
+                softmax_args = SoftmaxArgs(*softmax.args, **softmax.kwargs)  # type: ignore[arg-type, assignment]
+            elif softmax.target == torch.ops.aten._safe_softmax.default:
+                softmax_args = SafeSoftmaxArgs(*softmax.args, **softmax.kwargs)  # type: ignore[arg-type, assignment]
+            add, softmax_dim = (
                 softmax_args.input,
                 softmax_args.dim,
-                softmax_args.half_to_float,
             )
-            assert isinstance(add, torch.fx.Node), type(add)
-            assert isinstance(softmax_dim, int), type(softmax_dim)
-            assert isinstance(softmax_half_to_float, bool), type(softmax_half_to_float)
             softmax_shape = extract_shape(softmax)
             # TODO support other dimension
             if softmax_dim != -1 and softmax_dim != len(softmax_shape) - 1:
@@ -295,10 +295,16 @@ class RemoveRedundantReshapePattern3(PassBase):
             # Check condition
             reshape_2_input_shape = extract_shape(reshape_2_input)
             reshape_3_input_shape = extract_shape(reshape_3_input)
-            if reshape_2_input_shape != reshape_3_input_shape:
+            if not broadcastable(reshape_2_input_shape, reshape_3_input_shape):
                 continue
             reshape_1_shape = extract_shape(reshape_1)
-            if reshape_2_input_shape != reshape_1_shape:
+            if (
+                reshape_2_input_shape != reshape_1_shape
+                and reshape_3_input_shape != reshape_1_shape
+            ):
+                continue
+            # Make sure the softmax axis length is unchanged.
+            if softmax_shape[-1] != reshape_1_shape[-1]:
                 continue
             # Assume `aten.add` and `aten.softmax` have only one user.
             if len(add.users) != 1:
@@ -311,8 +317,7 @@ class RemoveRedundantReshapePattern3(PassBase):
             set_new_meta_val(add)
             # Update softmax
             if softmax_dim == len(softmax_shape) - 1:
-                updated_dim = len(extract_shape(reshape_2_input)) - 1
-                softmax.args = (add, updated_dim, softmax_half_to_float)
+                softmax.update_arg(1, -1)  # (index, last_dim)
             set_new_meta_val(softmax)
 
             reshape_1.replace_all_uses_with(softmax, propagate_meta=False)

tico/utils/utils.py

@@ -17,6 +17,7 @@ import subprocess
 import typing
 import warnings
 from functools import wraps
+from typing import List
 
 import torch
 from circle_schema import circle
@@ -341,3 +342,39 @@ def get_quant_dtype(qmin: int, qmax: int):
         return known_ranges[(qmin, qmax)]
     else:
         raise ValueError(f"Unsupported quantization range: ({qmin}, {qmax})")
+
+
+def broadcastable(
+    shape_a: List[int] | torch.Size, shape_b: List[int] | torch.Size
+) -> bool:
+    """
+    Return **True** if two shapes are broadcast-compatible under the standard
+    NumPy/PyTorch rules.
+
+    Broadcasting rule
+    --------------------------------
+    - Align the shapes **right-to-left**.
+    - For each aligned dimension `(a, b)` one of the following must hold
+      - `a == b`  (sizes match)
+      - `a == 1`  (shape-A can repeat along that dim)
+      - `b == 1`  (shape-B can repeat along that dim)
+    - When one shape is shorter, treat its missing leading dims as `1`.
+
+    Examples
+    --------
+    >>> _broadcastable([8, 16, 32], [16, 32])
+    True
+    >>> _broadcastable([8, 16, 32], [1, 32])
+    True
+    >>> _broadcastable([8, 16, 32], [8, 32, 16])
+    False
+    """
+    # Walk from the last dim to the front
+    len_a, len_b = len(shape_a), len(shape_b)
+    max_len = max(len_a, len_b)
+    for i in range(1, max_len + 1):
+        dim_a = shape_a[-i] if i <= len_a else 1
+        dim_b = shape_b[-i] if i <= len_b else 1
+        if dim_a != 1 and dim_b != 1 and dim_a != dim_b:
+            return False
+    return True

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

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

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

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=FjyTxaQjfrgfUCPfMzx5WnfpSY0Vbksu3gx8m6Tc5xY,1743
+tico/__init__.py,sha256=mEQHwb4j_KQrsDMO6gp2USb5RCYNhd0MOxL4P_EGMW8,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
@@ -82,12 +82,12 @@ tico/passes/lower_pow2_to_mul.py,sha256=imx9CoKG4bLyNqv4F-Z203s_P0-0SdRH-y4_Q0PT
 tico/passes/lower_to_resize_nearest_neighbor.py,sha256=4bIxPSyNEzznTw8f8D9hMNXbZ0KmMPPPfRvITXonEz0,8881
 tico/passes/lower_to_slice.py,sha256=6xK7A2hdqdBdN2XRyt2rFGnqJGcaXxmgndis4kn2q2w,7112
 tico/passes/merge_consecutive_cat.py,sha256=u1E-7axmX4yw0TMlFO5jXB0rn3xe5ASmqXcR42c9eNA,2722
-tico/passes/ops.py,sha256=BmUtfBPrSmCd2TsFc6PtOmhuznuFNBNHl12RS7CS0KI,2836
+tico/passes/ops.py,sha256=XzaKC_FpsfJLpnU4JlL9X-HVarWKm8cX0iiRgx9bMOs,2909
 tico/passes/remove_nop.py,sha256=5QE3inFsXgzyPT_t7pKeXNqD1LRf6ed_Mp7YMadA6AI,2707
 tico/passes/remove_redundant_assert_nodes.py,sha256=3a2xEQ2iPY7Gqg8jZi8G5bfDDrK2kOO1OHCMv_gJGz0,1592
 tico/passes/remove_redundant_expand.py,sha256=7st92AbWOl7yzM0Y5seaZJQKMFHqkYpH3qYMOlAU5lk,2234
-tico/passes/remove_redundant_permute.py,sha256=sS53eTY4sSnpZWDaaHN8czUmzNwmqh1lF90nYamXzac,3566
-tico/passes/remove_redundant_reshape.py,sha256=aPZcDR0kBExEsWCYfBbLulm_wcjJNnGjn4mgrUIPdpU,16810
+tico/passes/remove_redundant_permute.py,sha256=PIS-ag1EiSLlXJeRjcZoapbvM_4LyGtLBzunomfkAYE,4236
+tico/passes/remove_redundant_reshape.py,sha256=fDTMIrQFoRHYWrgGqEHgIg-azy4Z1Pod-Gfv7Gwplho,17010
 tico/passes/remove_redundant_slice.py,sha256=BAfSkA5jDIEhYx4nMnu6cJadQle3YTw5y39ZLiYfJJ8,2109
 tico/passes/remove_redundant_to_copy.py,sha256=uTIjAn3Eli_RvXC-QOqxBAkV_whDBkkNhu-mvNKAEhs,3136
 tico/passes/restore_linear.py,sha256=UMMHdLmRGq9bfJx_0L9lL2UQBd51PGNP0WywO8KdrDM,4066
@@ -189,15 +189,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=NAa3ZX5G-UCQwmz5WnFl0iCEra24PMY5wC0MyX7smUg,11156
+tico/utils/utils.py,sha256=7t4U3Bs7SV7y3PCdgE9bEAAAI2u1mzQEUP81oP_XHHU,12334
 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.dev250608.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250608.dist-info/METADATA,sha256=LJYAjDzOfuHcUGzJvc-Cgc_QKCnVUHFY2yr9IY8EkgU,8633
-tico-0.1.0.dev250608.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250608.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250608.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250608.dist-info/RECORD,,
+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,,