tico 0.1.0.dev250414__py3-none-any.whl → 0.1.0.dev250415__py3-none-any.whl

tico/__init__.py

@@ -22,7 +22,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.dev250414'
+__version__ = '0.1.0.dev250415'
 
 
 if Version(torch.__version__) < Version("2.5.0"):

tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py

@@ -32,6 +32,7 @@ from tico.utils.validate_args_kwargs import (
     LinearArgs,
     MulTensorArgs,
     PermuteArgs,
+    ReshapeArgs,
 )
 
 
@@ -102,6 +103,26 @@ class InsertQuantizeOnDtypeMismatch(PassBase):
     output dtype (int16) is updated to the input dtype (uint8), which breaks the semantics.
     This problem can occur in the tools (ex: circle2circle) that automatically apply type inference.
     - To resolve the issue, we insert quantize operators not to violate circle's type inference logic.
+    - NOTE For some cases, Quantize Op is inserted before the operators.
+
+    Let's assume Reshape Op's input is int16 and output is uint8. There are two possible places to insert
+    Quantize Op.
+
+    1. Insert Quantize before Reshape.
+
+    ```
+    Predecessor (int16)-> Quantize (uint8) -> Reshape (uint8) -> ...
+    ```
+
+    2. Insert Quantize after Reshape.
+
+    ```
+    Predecessor (int16)-> Reshape (int16) -> Quantize (uint8) -> ...
+    ```
+
+    Comparing 1) and 2), the difference is that Reshape operation is conducted in uint8 or int16.
+    We go with 1), which does Reshape in uint8, for faster execution. Note that Reshape Op does not
+    change the value, so its dytpe does not affect accuracy.
     """
 
     def __init__(self):
@@ -264,7 +285,38 @@ class InsertQuantizeOnDtypeMismatch(PassBase):
                     quantize = _insert_quantize_op_before(node, inp)
 
                     quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
-                    node.meta[QPARAM_KEY] = _u8_to_i16(node.meta[QPARAM_KEY])
+                    logger.debug(
+                        f"quantize_per_tensor.default is inserted before {node.name}."
+                    )
+                elif qparam_dtype(inp) == "uint8" and qparam_dtype(node) == "int16":
+                    quantize = _insert_quantize_op_after(node)
+
+                    quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
+                    node.meta[QPARAM_KEY] = _i16_to_u8(node.meta[QPARAM_KEY])
+                    logger.debug(
+                        f"quantize_per_tensor.default is inserted after {node.name}."
+                    )
+                else:
+                    raise NotYetSupportedError("Unsupported dtype")
+            elif node.target == torch.ops.aten.reshape.default:
+                reshape_args = ReshapeArgs(*node.args, **node.kwargs)
+                inp = reshape_args.input
+
+                if QPARAM_KEY not in inp.meta:
+                    continue
+
+                if QPARAM_KEY not in node.meta:
+                    continue
+
+                if qparam_dtype(inp) == qparam_dtype(node):
+                    continue
+
+                if qparam_dtype(inp) == "int16" and qparam_dtype(node) == "uint8":
+                    # A new Quantize Op (s16 to u8) is inserted before (not after)
+                    # reshape Op to reduce tensor size ealier
+                    quantize = _insert_quantize_op_before(node, inp)
+
+                    quantize.meta[QPARAM_KEY] = copy.deepcopy(node.meta[QPARAM_KEY])
                     logger.debug(
                         f"quantize_per_tensor.default is inserted before {node.name}."
                     )

tico/passes/lower_to_resize_nearest_neighbor.py

@@ -35,32 +35,6 @@ class LowerToResizeNearestNeighbor(PassBase):
     This pass lowers `aten.index` and `aten.upsample_nearest2d.vec` to `circle_custom.resize_nearest_neighbor` when it is possible.
 
     Until torch 2.7, `torch.nn.functional.interpolate` is converted to `aten.index` op.
-
-        [EXAMPLE]
-        class InterpolateDouble(torch.nn.Module):
-            def __init__(self):
-                super().__init__()
-
-            def forward(self, x):
-                return torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest")
-
-            def get_example_inputs(self):
-                return (torch.randn(1, 2, 3, 4),)
-
-        [EXPORTED GRAPH]
-        [constants]
-            _prop_tensor_constant0 = tensor([0, 0, 1, 1, 2, 2, 3, 3]
-            _prop_tensor_constant1 = tensor([[0], [0], [1], [1], [2], [2]])
-
-        [graph]
-            %_prop_tensor_constant0 : [num_users=1] = placeholder[target=_prop_tensor_constant0]
-            %_prop_tensor_constant1 : [num_users=1] = placeholder[target=_prop_tensor_constant1]
-            %x : [num_users=1] = placeholder[target=x]
-            %_to_copy : [num_users=1] = call_function[target=torch.ops.aten._to_copy.default](args = (%x,), kwargs = {dtype: torch.float32})
-            %index : [num_users=1] = call_function[target=torch.ops.aten.index.Tensor](args = (%_to_copy, [None, None, %_prop_tensor_constant1, %_prop_tensor_constant0]), kwargs = {})
-            %_to_copy_3 : [num_users=1] = call_function[target=torch.ops.aten._to_copy.default](args = (%index,), kwargs = {dtype: torch.float32})
-            return (_to_copy_3,)
-
         [BEFORE PASS]
         input - aten.index - output
 
@@ -68,6 +42,11 @@ class LowerToResizeNearestNeighbor(PassBase):
         input - aten.permute(NCHW_to_NHWC) - circle_custom.resize_nearest_neighbor - aten.permute(NHWC_to_NCHW) - output
 
     Since torch 2.8, `torch.nn.functional.interpolate` is converted to aten.upsample_nearest2d.vec` op.
+        [BEFORE PASS]
+        input - aten.upsample_nearest2d.vec - output
+
+        [AFTER PASS]
+        input - aten.permute(NCHW_to_NHWC) - circle_custom.resize_nearest_neighbor - aten.permute(NHWC_to_NCHW) - output
     """
 
     def __init__(self):

tico/passes/lower_to_slice.py

@@ -17,37 +17,56 @@ 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 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 SelectCopyIntArgs
+from tico.utils.validate_args_kwargs import IndexSelectArgs, SelectCopyIntArgs
 
 
-@trace_const_diff_on_pass
-class LowerToSlice(PassBase):
+def passes():
     """
     This pass lowers aten.ops.select/selct_copy.int to aten.ops.slice.
     We support only when it is index in args, which is a constant tensor.
     Since the index in node'args isn't constant tensor, we can't support converting the below op list yet.
-    - torch.ops.aten.index_select.default
+
+    TODO Support below with const indices
     - torch.ops.aten.embedding.default
     - torch.ops.aten.index.Tensor
+    """
+    return [
+        LowerSelectCopyToSlice(),
+        LowerIndexSelectToSlice(),
+    ]
+
 
+@trace_const_diff_on_pass
+class LowerSelectCopyToSlice(PassBase):
+    """
     [before]
-            input (tensor, dim, *index)
+            input
                 |
-            select
+            select (tensor, dim, *index)
                 |
             output
 
     [after]
 
-            input (tensor, dim, *index)
+            input
                 |
             slice (input=tensor, dim=dim, start=index, end=index+1, step=1)
                 |
@@ -110,3 +129,98 @@ class LowerToSlice(PassBase):
         graph_module.recompile()
 
         return PassResult(modified)
+
+
+@trace_const_diff_on_pass
+class LowerIndexSelectToSlice(PassBase):
+    """
+
+    [before]
+            input
+                |
+            index_select.default  (tensor, dim, *index)
+                |
+            output
+
+    [after]
+
+            input
+                |
+            slice (input=tensor, dim=dim, start=index, end=index+1, step=1)
+                |
+            reshape (input=slice_copy, size=select_shape)
+                |
+            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
+
+            input_shape = extract_shape(input)
+            if dim < 0:
+                dim = dim % len(input_shape)
+
+            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 not is_single_value_tensor(index):
+                # need to be lowered by LowerIndexSelect pass
+                continue
+            index_int = index.item()  # convert scalar tensor to int
+
+            start = index_int
+            end = index_int + 1
+            step = 1
+            slice_copy_args = (input, dim, start, end, step)
+
+            with graph.inserting_after(node):
+                # slice
+                slice_node = graph.call_function(
+                    torch.ops.aten.slice.Tensor, args=slice_copy_args
+                )
+                node_shape = extract_shape(node)
+            with graph.inserting_after(slice_node):
+                # reshape
+                reshape_args = (slice_node, list(node_shape))
+                reshape_node = graph.call_function(
+                    torch.ops.aten.reshape.default, args=reshape_args
+                )
+                node.replace_all_uses_with(reshape_node, propagate_meta=False)
+
+            modified = True
+            logger.debug(
+                f"{node.name} is replaced with {slice_node.name} and {reshape_node.name} operators"
+            )
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)

tico/utils/convert.py

@@ -58,7 +58,7 @@ from tico.passes.legalize_predefined_layout_operators import (
 )
 from tico.passes.lower_pow2_to_mul import LowerPow2ToMul
 from tico.passes.lower_to_resize_nearest_neighbor import LowerToResizeNearestNeighbor
-from tico.passes.lower_to_slice import LowerToSlice
+from tico.passes.lower_to_slice import passes as LowerToSlicePasses
 from tico.passes.merge_consecutive_cat import MergeConsecutiveCat
 from tico.passes.remove_nop import RemoveNop
 from tico.passes.remove_redundant_assert_nodes import RemoveRedundantAssertionNodes
@@ -224,7 +224,7 @@ def convert_exported_module_to_circle(
             LegalizePreDefinedLayoutOperators(),
             LowerPow2ToMul(),
             ConvertConv1dToConv2d(),
-            LowerToSlice(),
+            *LowerToSlicePasses(),
         ]
     )
     circle_legalize.run(exported_program)

tico/utils/validate_args_kwargs.py

@@ -852,7 +852,7 @@ class ReshapeArgs:
 @dataclass
 class ResizeNearestNeighborArgs:
     """
-    # Maps from `torch.nn.functional.interpolate(x, scale_factor=scale_factor, mode='nearest')` case.
+    # Mapped from `torch.nn.functional.interpolate(x, scale_factor=scale_factor, mode='nearest')` case.
     """
 
     input: torch.fx.Node

{tico-0.1.0.dev250414.dist-info → tico-0.1.0.dev250415.dist-info}/LICENSE

@@ -1,6 +1,6 @@
 This file provides full text of licenses used in this project
 
-- Apache Licence 2.0
+- Apache License 2.0
 - BSD 3-Clause
 
 ...............................................................................

{tico-0.1.0.dev250414.dist-info → tico-0.1.0.dev250415.dist-info}/METADATA

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

{tico-0.1.0.dev250414.dist-info → tico-0.1.0.dev250415.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=drwOk40YiN-xglot_aW8vG4uxq2GWx4F853Gxd6Jdv0,1181
+tico/__init__.py,sha256=2PwjELyAXeZkJt02ODDKWpNPwQj5W4GtKFjH-VaZBdk,1181
 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
@@ -51,7 +51,7 @@ tico/experimental/quantization/evaluation/executor/circle_executor.py,sha256=eCC
 tico/experimental/quantization/evaluation/executor/triv24_executor.py,sha256=sUoXl6oOO2arAKaNjOBg7HiQja145_Jv6qgY7XtR7A8,5159
 tico/experimental/quantization/passes/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/experimental/quantization/passes/fold_quant_ops.py,sha256=Jq5wmQDhdjsXxae2p6TnZj2gY5UMBEQ-sHkTodgkfUs,3327
-tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py,sha256=pqKPk-F5AYsas3aiCC0SmNginvHFcPKjSgmOg4JWlJ0,10822
+tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py,sha256=i4rkM1vlN85fXA9oOrU25o8KWAaqA65NKngTX6MgctQ,12960
 tico/experimental/quantization/passes/propagate_qparam_backward.py,sha256=TGtyW0Z2qOTgVIasBdGRgbwH31YYd6ek7OvLTmCV614,3118
 tico/experimental/quantization/passes/propagate_qparam_forward.py,sha256=RhUHGCR2RpBO5KYkQ7Z8U5u7HEwDq2wdKHLKAJCi-5c,5138
 tico/experimental/quantization/passes/remove_weight_dequant_op.py,sha256=V-0tZqLxHbB0Dvw2D69aesfGdr3vmhvFlfNAIYnGTuo,6361
@@ -79,8 +79,8 @@ tico/passes/fuse_redundant_reshape_to_mean.py,sha256=SzHLLL2Yfj4k1c2L5i4PVI9EFUi
 tico/passes/legalize_causal_mask_value.py,sha256=KZc_UPk7CGPXO35JOu6dVrOzRJx-ZpJoVvuzz-GvQek,4080
 tico/passes/legalize_predefined_layout_operators.py,sha256=OKVKxw039Bl9df7YCt17mWOHaGhQLWkJeSIPQpWdNBM,16224
 tico/passes/lower_pow2_to_mul.py,sha256=imx9CoKG4bLyNqv4F-Z203s_P0-0SdRH-y4_Q0PTZVo,2304
-tico/passes/lower_to_resize_nearest_neighbor.py,sha256=EXH-xuXggLFSLaDEknxG_TdY4wHznP4zDG6M3qzq1A0,9982
-tico/passes/lower_to_slice.py,sha256=XUZ5DLDmWQeR2LU7yC4DDWK0lDwsdCS7BDiqambncA4,3651
+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/remove_nop.py,sha256=5QE3inFsXgzyPT_t7pKeXNqD1LRf6ed_Mp7YMadA6AI,2707
@@ -175,7 +175,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=wQrcrnZ942_4SfhEW_7hIVi0tncGor-o7zuf2kmk9Io,1803
 tico/utils/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
-tico/utils/convert.py,sha256=bu-9T_tBL9UawITqMPDo-T2WkI-d64XZ13NHXffKIGk,11661
+tico/utils/convert.py,sha256=KCllPnvQ8bjEYR1yI72s9aNBp7Py1CzIEEpYSYZcu60,11684
 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
@@ -187,10 +187,10 @@ tico/utils/passes.py,sha256=kGmDe__5cPaO6i5EDAoXSVe6yXEoX9hAny4ROb3ZEmQ,2409
 tico/utils/register_custom_op.py,sha256=FbMcrg8o5vKWC_aoVxL2GrIcR14KFi1yKG0mFGqXkPY,21595
 tico/utils/trace_decorators.py,sha256=ddLIiKQfSaQrxgF1kNpwjFTQnXENzeSfcr1kuAW4jGI,3221
 tico/utils/utils.py,sha256=pybDU1LoNhjEplANig11lboX9yzYRkvFCSmyYth_2Do,10359
-tico/utils/validate_args_kwargs.py,sha256=DXW7W5x-6pQR43_q4EFoeWE5pn1Nm7gNsrzaO5avm4k,24724
-tico-0.1.0.dev250414.dist-info/LICENSE,sha256=dAq6L2M49W6wEPmkyVtYxE3Omm3oInXPd6qw5mQX6wY,12644
-tico-0.1.0.dev250414.dist-info/METADATA,sha256=NxIJgt_iKq8KOz07g5VLVRyCv7Q9UQXJYML-EWcQYoY,7353
-tico-0.1.0.dev250414.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250414.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250414.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250414.dist-info/RECORD,,
+tico/utils/validate_args_kwargs.py,sha256=krT68b5CfBI9rxBIOsgYSy0LfEJqLfKfRikkp8ep9oQ,24726
+tico-0.1.0.dev250415.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250415.dist-info/METADATA,sha256=zqEFm2T9WSrEat4XqjQXesrRZTWeXN5BoaOZrdD6mhI,7353
+tico-0.1.0.dev250415.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250415.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250415.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250415.dist-info/RECORD,,