tico 0.1.0.dev250714__py3-none-any.whl → 0.1.0.dev250716__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.dev250714"
+__version__ = "0.1.0.dev250716"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/passes/cast_clamp_mixed_type_args.py

@@ -0,0 +1,169 @@
+# 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 import ExportedProgram
+
+from tico.passes import ops
+
+from tico.serialize.circle_mapping import extract_torch_dtype
+from tico.utils import logging
+from tico.utils.graph import create_node
+from tico.utils.passes import PassBase, PassResult
+from tico.utils.trace_decorators import trace_graph_diff_on_pass
+from tico.utils.utils import is_target_node, set_new_meta_val
+from tico.utils.validate_args_kwargs import ClampArgs
+
+
+@trace_graph_diff_on_pass
+class CastClampMixedTypeArgs(PassBase):
+    """
+    This pass ensures consistent dtypes for clamp operations by:
+    1. Converting min/max arguments to match output dtype when provided
+    2. Inserting cast operations when input dtype differs from output dtype
+
+    Behavior Examples:
+    - When input dtype differs from output:
+        Inserts _to_copy operation to convert input
+    - When min/max dtype differs from output:
+        Converts min/max values to output dtype
+
+    (Case 1, if input dtype is different from output dtype)
+    [before]
+
+            input               min(or max)
+           (dtype=int)         (dtype=float)
+              |                    |
+            clamp <----------------+
+              |
+            output
+           (dtype=float)
+
+    [after]
+
+            input             min(or max)
+           (dtype=int)       (dtype=float)
+              |                  |
+            cast                 |
+          (in=int, out=float)    |
+              |                  |
+            clamp <--------------+
+              |
+            output
+           (dtype=float)
+
+    (Case 2, if min(or max) dtype is different from output dtype)
+    [before]
+
+            input               min(or max)
+           (dtype=float)       (dtype=int)
+              |                    |
+            clamp <----------------+
+              |
+            output
+           (dtype=float)
+
+    [after]
+
+            input             min(or max)
+           (dtype=float)     (dtype=float)
+              |                  |
+            clamp <--------------+
+              |
+            output
+           (dtype=float)
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def convert(self, exported_program: ExportedProgram, node: torch.fx.Node) -> bool:
+        logger = logging.getLogger(__name__)
+        modified = False
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+
+        # clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor
+        args = ClampArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+
+        input = args.input
+        min = args.min
+        max = args.max
+
+        input_dtype = extract_torch_dtype(input)
+        output_dtype = extract_torch_dtype(node)
+
+        def _convert_arg(arg, arg_name: str):
+            if arg is None:
+                return False
+
+            arg_dtype = torch.tensor(arg).dtype
+            arg_idx = node.args.index(arg)
+            if arg_dtype != output_dtype:
+                assert output_dtype in [torch.float, torch.int]
+                if output_dtype == torch.float:
+                    arg = float(arg)
+                else:
+                    arg = int(arg)
+                node.update_arg(arg_idx, arg)
+                logger.debug(
+                    f"Casting {arg_name} value from {arg_dtype} to {output_dtype} for clamp operation at {node.name}"
+                )
+                return True
+            return False
+
+        modified |= _convert_arg(min, "min")
+        modified |= _convert_arg(max, "max")
+
+        if input_dtype != output_dtype:
+            logger.debug(
+                f"Inserting cast from {input_dtype} to {output_dtype} for input {input.name}"
+            )
+            with graph.inserting_after(input):
+                to_copy = create_node(
+                    graph,
+                    torch.ops.aten._to_copy.default,
+                    (input,),
+                    {"dtype": output_dtype},
+                    origin=input,
+                )
+                set_new_meta_val(to_copy)
+                node.update_arg(node.args.index(input), to_copy)
+
+            modified = True
+
+        return modified
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        target_op = ops.aten.clamp
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        modified = False
+        for node in graph.nodes:
+            if not is_target_node(node, target_op):
+                continue
+
+            modified |= self.convert(exported_program, node)
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)

tico/passes/ops.py

@@ -73,6 +73,7 @@ class AtenOps:
             torch.ops.aten.view.default,
             torch.ops.aten.view_copy.default,
         ]
+        self._to_copy = [torch.ops.aten._to_copy.default]
 
 
 aten = AtenOps()

tico/utils/convert.py

@@ -35,6 +35,7 @@ from tico.experimental.quantization.passes.remove_weight_dequant_op import (
     RemoveWeightDequantOp,
 )
 from tico.passes.cast_aten_where_arg_type import CastATenWhereArgType
+from tico.passes.cast_clamp_mixed_type_args import CastClampMixedTypeArgs
 from tico.passes.cast_mixed_type_args import CastMixedTypeArgs
 from tico.passes.const_prop_pass import ConstPropPass
 from tico.passes.convert_conv1d_to_conv2d import ConvertConv1dToConv2d
@@ -251,6 +252,7 @@ def convert_exported_module_to_circle(
             ConvertConv1dToConv2d(),
             *LowerToSlicePasses(),
             FuseLeadingUnsqueezeReshape(),
+            CastClampMixedTypeArgs(),
         ]
     )
     circle_legalize.run(exported_program)

tico/utils/model.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+from pathlib import Path
 from typing import Any
 
 from tico.interpreter import infer
@@ -32,6 +33,6 @@ class CircleModel:
             buf = bytes(f.read())
         return CircleModel(buf)
 
-    def save(self, circle_path: str) -> None:
+    def save(self, circle_path: str | Path) -> None:
         with open(circle_path, "wb") as f:
             f.write(self.circle_binary)

tico/utils/pytree_utils.py

@@ -0,0 +1,134 @@
+import threading
+
+import torch
+from packaging.version import Version
+
+from tico.utils import logging
+from tico.utils.installed_packages import is_transformers_installed
+
+__all__ = ["register_dynamic_cache"]
+
+
+def register_dynamic_cache():
+    PyTreeRegistryHelper().register_dynamic_cache()
+
+
+class PyTreeRegistryHelper:
+    """
+    Thread-safe singleton helper class for registering custom PyTree nodes.
+
+    This class provides functionality to register DynamicCache as a PyTree node
+    for torch.export compatibility. This registration is only needed for
+    transformers versions below 4.50.0.
+
+    Thread Safety:
+    - Uses a class-level threading.Lock() to ensure thread-safe singleton instantiation
+    - Uses the same lock to protect the registration process from concurrent calls
+    """
+
+    _instance = None  # Class variable to hold the singleton instance
+    _has_called = False  # Flag to track if registration has been performed
+    _lock = threading.Lock()  # Class-level lock for thread-safe operations
+
+    def __init__(self):
+        """Private constructor to prevent direct instantiation"""
+        pass
+
+    def __new__(cls, *args, **kwargs):
+        """
+        Thread-safe singleton instance creation using double-checked locking pattern.
+
+        Returns:
+            PyTreeRegistryHelper: The singleton instance of this class
+        """
+        if not cls._instance:
+            with cls._lock:  # Acquire lock for thread-safe instantiation
+                if not cls._instance:  # Double-check after acquiring lock
+                    cls._instance = super().__new__(cls)
+        return cls._instance
+
+    def register_dynamic_cache(self):
+        """
+        Registers DynamicCache as a PyTree node for torch.export compatibility.
+
+        This method is thread-safe and idempotent - it will only perform the
+        registration once, even if called multiple times from different threads.
+
+        Note:
+            This registration is only needed for transformers versions below 4.50.0.
+
+        Raises:
+            ImportError: If transformers package is not installed
+        """
+        with self._lock:  # Acquire lock for thread-safe registration
+            if self.__class__._has_called:
+                logger = logging.getLogger(__name__)
+                logger.debug("register_dynamic_cache already called, skipping")
+                return
+
+            self.__class__._has_called = True
+            logger = logging.getLogger(__name__)
+            logger.info("Registering DynamicCache PyTree node")
+
+        if not is_transformers_installed:  # type: ignore[truthy-function]
+            raise ImportError("transformers package is not installed")
+
+        import transformers
+
+        HAS_TRANSFORMERS_LESS_4_50_0 = Version(transformers.__version__) < Version(
+            "4.50.0"
+        )
+        if not HAS_TRANSFORMERS_LESS_4_50_0:
+            return
+
+        from transformers.cache_utils import DynamicCache
+
+        def _flatten_dynamic_cache(dynamic_cache: DynamicCache):
+            if not isinstance(dynamic_cache, DynamicCache):
+                raise RuntimeError(
+                    "This pytree flattening function should only be applied to DynamicCache"
+                )
+            HAS_TORCH_2_6_0 = Version(torch.__version__) >= Version("2.6.0")
+            if not HAS_TORCH_2_6_0:
+                logger = logging.getLogger(__name__)
+                logger.warning_once(
+                    "DynamicCache + torch.export is tested on torch 2.6.0+ and may not work on earlier versions."
+                )
+            dictionary = {
+                "key_cache": getattr(dynamic_cache, "key_cache"),
+                "value_cache": getattr(dynamic_cache, "value_cache"),
+            }
+            return torch.utils._pytree._dict_flatten(dictionary)
+
+        def _flatten_with_keys_dynamic_cache(dynamic_cache: DynamicCache):
+            dictionary = {
+                "key_cache": getattr(dynamic_cache, "key_cache"),
+                "value_cache": getattr(dynamic_cache, "value_cache"),
+            }
+            return torch.utils._pytree._dict_flatten_with_keys(dictionary)
+
+        def _unflatten_dynamic_cache(values, context: torch.utils._pytree.Context):
+            dictionary = torch.utils._pytree._dict_unflatten(values, context)
+            cache = DynamicCache()
+            for k, v in dictionary.items():
+                setattr(cache, k, v)
+            return cache
+
+        def _flatten_dynamic_cache_for_fx(cache, spec):
+            dictionary = {
+                "key_cache": getattr(cache, "key_cache"),
+                "value_cache": getattr(cache, "value_cache"),
+            }
+            return torch.fx._pytree._dict_flatten_spec(dictionary, spec)
+
+        torch.utils._pytree.register_pytree_node(
+            DynamicCache,
+            _flatten_dynamic_cache,
+            _unflatten_dynamic_cache,
+            serialized_type_name=f"{DynamicCache.__module__}.{DynamicCache.__name__}",
+            flatten_with_keys_fn=_flatten_with_keys_dynamic_cache,
+        )
+        # TODO: This won't be needed in torch 2.7+.
+        torch.fx._pytree.register_pytree_flatten_spec(
+            DynamicCache, _flatten_dynamic_cache_for_fx
+        )

{tico-0.1.0.dev250714.dist-info → tico-0.1.0.dev250716.dist-info}/METADATA

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

{tico-0.1.0.dev250714.dist-info → tico-0.1.0.dev250716.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tico/__init__.py,sha256=xbMeC4-9n5Z8Lu1B6yqbmE10G6wmHdhpue7kD7skYtM,1743
+tico/__init__.py,sha256=7Bu_kNw98Z5aDMnphDfqKQ00FlgKKT8iym416i5oNBI,1743
 tico/pt2_to_circle.py,sha256=gu3MD4Iqc0zMZcCZ2IT8oGbyj21CTSbT3Rgd9s2B_9A,2767
 tico/config/__init__.py,sha256=xZzCXjZ84qE-CsBi-dfaL05bqpQ3stKKfTXhnrJRyVs,142
 tico/config/base.py,sha256=anwOiJFkUxUi7Cef573JgQcjk6S-FSi6O_TLjYASW-g,1244
@@ -61,6 +61,7 @@ tico/interpreter/infer.py,sha256=1ZFe3DVMR2mlwBosoedqoL0-CGN_01CKLgMgxuw62KA,486
 tico/interpreter/interpreter.py,sha256=tGbluCbrehTCqBu8mtGDNzby_ieJ2ry8_RH_eC0CQxk,3828
 tico/passes/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
 tico/passes/cast_aten_where_arg_type.py,sha256=ybtGj1L7_2zGyfb_G-_y1N1mRgKHVq6fBZc-9-fH9sA,7229
+tico/passes/cast_clamp_mixed_type_args.py,sha256=m3_HpXLywWmWERfE5lM5PgvjBod7C4BWu_Q-TkRyO8k,5387
 tico/passes/cast_mixed_type_args.py,sha256=ArJpIPnQP1LPNaWIwee13Nbj749_awFKDO-pAYvdYvI,7618
 tico/passes/const_prop_pass.py,sha256=QOeR2u3fo9ZhWXRhfAUW1dTtuWgqgoqdDJoQ516UDbQ,11532
 tico/passes/convert_conv1d_to_conv2d.py,sha256=7YljWJQBX5vBUMgGgRv8TvbJ9UpEL9hf4ZU3dNUhEZ8,5301
@@ -84,7 +85,7 @@ tico/passes/lower_pow2_to_mul.py,sha256=nfJXa9ZTZMiLg6ownSyvkM4KF2z9tZW34Q3CCWI_
 tico/passes/lower_to_resize_nearest_neighbor.py,sha256=N6F56Of8Aiv-KIiYLHnh33WX72W60ZVQSBEYWHdYqNQ,9005
 tico/passes/lower_to_slice.py,sha256=0qAX3WzZdyMFDW4DiO9b5JFXd4rL1-0doBT6lJvaw_I,7260
 tico/passes/merge_consecutive_cat.py,sha256=BYmiU170DsrHQMj7gMe7U6ZpndrX-S4OpvJweDdspec,2701
-tico/passes/ops.py,sha256=XzaKC_FpsfJLpnU4JlL9X-HVarWKm8cX0iiRgx9bMOs,2909
+tico/passes/ops.py,sha256=cSj3Sk2x2cOE9b8oU5pmSa_rHr-iX2lORzu3N_UHMSQ,2967
 tico/passes/remove_nop.py,sha256=Hf91p_EJAOC6DyWNthash0_UWtEcNc_M7znamQfYQ5Y,2686
 tico/passes/remove_redundant_assert_nodes.py,sha256=IONd3xBy6I8tH6_Y1eN3_eCHH7WTC8soBgjXzOju9cQ,1612
 tico/passes/remove_redundant_expand.py,sha256=5SIqN7eIIcqF68tlrB31n1482jSBSBOgKb1wddLX6lw,2197
@@ -182,16 +183,17 @@ tico/serialize/operators/op_view.py,sha256=5EMww-ve17Vm9XPuV03Tn7vJsjpU2J8U4d_FO
 tico/serialize/operators/op_where.py,sha256=doE81GSwygrPBm3JIfN9w7kKXxeIYKxgk0eoY22QIcg,2845
 tico/serialize/operators/utils.py,sha256=lXGpEJW1h8U_-gfc6EWjvvSiq3yJ9P-v1v3EMRT_pSk,2954
 tico/utils/__init__.py,sha256=IO6FP_xYbGy0dW0HL26GXD3ouxARaxCK7bz9dn4blPQ,26
-tico/utils/convert.py,sha256=11Ps0i4-3Fmcts_PZO5Eo8rRL7FcLV33eHz_G97MnCg,12865
+tico/utils/convert.py,sha256=w4l7fnqbiVACOU5-OXr8Ebyl4EMeeBz6vwUSuOS_CtI,12977
 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
 tico/utils/graph.py,sha256=Y6aODsnc_-9l61oanknb7K1jqJ8B35iPypOKkM0Qkk0,9149
 tico/utils/installed_packages.py,sha256=J0FTwnkCGs0MxRWoCMYAqiwH7Z0GWFDLV--x-IndSp4,1017
 tico/utils/logging.py,sha256=IlbBWscsaHidI0dNqro1HEXAbIcbkR3BD5ukLy2m95k,1286
-tico/utils/model.py,sha256=Uqc92AnJXQ2pbvctS2z2F3Ku3yNrwXZ9O33hZVis7is,1250
+tico/utils/model.py,sha256=pPOIjD0qjQirLibiRxxfjOR6efimOcDAd9R-74eus-k,1282
 tico/utils/padding.py,sha256=jyNhGmlLZfruWZ6n5hll8RZOFg85iCZP8OJqnHGS97g,3293
 tico/utils/passes.py,sha256=kGmDe__5cPaO6i5EDAoXSVe6yXEoX9hAny4ROb3ZEmQ,2409
+tico/utils/pytree_utils.py,sha256=jrk3N6X6LiUnBCX_gM1K9nywbVAJBVnszlTAgeIeDUc,5219
 tico/utils/register_custom_op.py,sha256=3-Yl6iYmx1qQA2igNHt4hYhQhQMkdPb7gF50LIY8yvc,27350
 tico/utils/serialize.py,sha256=AQXMBOLu-Kg2Rn-qbqsAtHndjZAZIavlKA0QFgJREHM,1420
 tico/utils/trace_decorators.py,sha256=ddLIiKQfSaQrxgF1kNpwjFTQnXENzeSfcr1kuAW4jGI,3221
@@ -201,9 +203,9 @@ 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.dev250714.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
-tico-0.1.0.dev250714.dist-info/METADATA,sha256=5RPA0JUl-L5MOz1_Ix5-8SZSpn628DHjCpNZHPvrO68,8430
-tico-0.1.0.dev250714.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
-tico-0.1.0.dev250714.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
-tico-0.1.0.dev250714.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
-tico-0.1.0.dev250714.dist-info/RECORD,,
+tico-0.1.0.dev250716.dist-info/LICENSE,sha256=kp4JLII7bzRhPb0CPD5XTDZMh22BQ7h3k3B7t8TiSbw,12644
+tico-0.1.0.dev250716.dist-info/METADATA,sha256=ol8z8iSa9e8WfJcDpebI1h_hXSQdi3jXwIOzIp-wAO4,8430
+tico-0.1.0.dev250716.dist-info/WHEEL,sha256=G16H4A3IeoQmnOrYV4ueZGKSjhipXx8zc8nu9FGlvMA,92
+tico-0.1.0.dev250716.dist-info/entry_points.txt,sha256=kBKYSS_IYrSXmUYevmmepqIVPScq5vF8ulQRu3I_Zf0,59
+tico-0.1.0.dev250716.dist-info/top_level.txt,sha256=oqs7UPoNSKZEwqsX8B-KAWdQwfAa7i60pbxW_Jk7P3w,5
+tico-0.1.0.dev250716.dist-info/RECORD,,