tico 0.1.0.dev250714__py3-none-any.whl → 0.1.0.dev251102__py3-none-any.whl

tico/__init__.py

@@ -20,8 +20,16 @@ from packaging.version import Version
 from tico.config import CompileConfigV1, get_default_config
 from tico.utils.convert import convert, convert_from_exported_program, convert_from_pt2
 
+__all__ = [
+    "CompileConfigV1",
+    "get_default_config",
+    "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.dev251102"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/config/base.py

@@ -31,7 +31,7 @@ class CompileConfigBase:
         config = cls()
         for key in config_dict:
             if key in config.to_dict():
-                assert type(config.get(key)) == bool
+                assert isinstance(config.get(key), bool)
                 config.set(key, config_dict[key])
 
         return config

tico/config/v1.py

@@ -20,6 +20,11 @@ from tico.config.base import CompileConfigBase
 @dataclass
 class CompileConfigV1(CompileConfigBase):
     legalize_causal_mask_value: bool = False
+    remove_constant_input: bool = False
+    convert_lhs_const_mm_to_fc: bool = False
+    convert_rhs_const_mm_to_fc: bool = True
+    convert_single_batch_lhs_const_bmm_to_fc: bool = False
+    convert_expand_to_slice_cat: bool = False
 
     def get(self, name: str):
         return super().get(name)

tico/passes/cast_aten_where_arg_type.py

@@ -176,7 +176,7 @@ class CastATenWhereArgType(PassBase):
             node_dtype = extract_torch_dtype(node)
             assert (
                 node_dtype == node_dtype_ori
-            ), f"Type casting doesn't change node's dtype."
+            ), "Type casting doesn't change node's dtype."
 
             logger.debug(
                 f"{to_cast.name}'s dtype was casted from {buf_data.dtype} to {dtype_to_cast}"

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/cast_mixed_type_args.py

@@ -41,6 +41,8 @@ ops_to_promote = {
     torch.ops.aten.ge.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
     torch.ops.aten.gt.Scalar: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
     torch.ops.aten.gt.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
+    torch.ops.aten.le.Scalar: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
+    torch.ops.aten.le.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
     torch.ops.aten.mul.Tensor: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
     torch.ops.aten.minimum.default: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
     torch.ops.aten.ne.Scalar: ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT,
@@ -124,7 +126,7 @@ class CastMixedTypeArgs(PassBase):
             if rhs_val.dtype == type_to_promote:
                 ori_type = lhs_val.dtype
                 arg_to_promote = lhs
-            assert arg_to_promote != None
+            assert arg_to_promote is not None
 
             if isinstance(arg_to_promote, torch.fx.Node):
                 with graph.inserting_after(arg_to_promote):
@@ -178,7 +180,7 @@ class CastMixedTypeArgs(PassBase):
             node_dtype = extract_torch_dtype(node)
             assert (
                 node_dtype == node_dtype_ori
-            ), f"Type casting doesn't change node's dtype."
+            ), "Type casting doesn't change node's dtype."
 
         graph.eliminate_dead_code()
         graph.lint()

tico/passes/const_prop_pass.py

@@ -301,7 +301,7 @@ class ConstPropPass(PassBase):
         graph.eliminate_dead_code()
         graph_module.recompile()
 
-        logger.debug(f"Constant nodes are propagated")
+        logger.debug("Constant nodes are propagated")
         # Constant folding can be done with only one time run. Let's set `modified` to False.
         modified = False
         return PassResult(modified)

tico/passes/convert_conv1d_to_conv2d.py

@@ -19,7 +19,7 @@ if TYPE_CHECKING:
 import torch
 from torch.export import ExportedProgram
 
-from tico.serialize.circle_graph import extract_shape
+from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 from tico.utils.graph import create_node

tico/passes/convert_expand_to_slice_cat.py

@@ -0,0 +1,153 @@
+# 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_shape
+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
+from tico.utils.validate_args_kwargs import ExpandArgs, ReshapeArgs
+
+
+@trace_graph_diff_on_pass
+class ConvertExpandToSliceCat(PassBase):
+    """
+    This pass replaces `aten.reshape` + `aten.expand` pattern by rewriting it using
+    a series of `aten.slice` and `aten.cat` operations.
+
+    This pass is specialized for expand of KVCache.
+    - Expects (batch, num_key_value_heads, seq_len, head_dim) as input shape of reshape
+    """
+
+    def __init__(self, enabled: bool = False):
+        super().__init__()
+        self.enabled = enabled
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        if not self.enabled:
+            return PassResult(False)
+
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+        modified = False
+
+        # This pass handles expand on EXPAND_DIM only
+        CAT_DIM = 1
+        EXPAND_DIM = 2
+
+        for node in graph.nodes:
+            if not isinstance(node, torch.fx.Node) or not is_target_node(
+                node, ops.aten.reshape
+            ):
+                continue
+
+            post_reshape = node
+            post_reshape_args = ReshapeArgs(*post_reshape.args, **post_reshape.kwargs)
+            post_reshape_input = post_reshape_args.input
+
+            if not isinstance(post_reshape_input, torch.fx.Node) or not is_target_node(
+                post_reshape_input, ops.aten.expand
+            ):
+                continue
+
+            expand = post_reshape_input
+            expand_args = ExpandArgs(*expand.args, **expand.kwargs)
+            expand_input = expand_args.input
+            expand_shape = extract_shape(expand)
+
+            if not isinstance(expand_input, torch.fx.Node) or not is_target_node(
+                expand_input, ops.aten.reshape
+            ):
+                continue
+
+            pre_reshape = expand_input
+            pre_reshape_args = ReshapeArgs(*pre_reshape.args, **pre_reshape.kwargs)
+            pre_reshape_input = pre_reshape_args.input
+            pre_reshape_shape = extract_shape(pre_reshape)
+
+            if pre_reshape_shape[EXPAND_DIM] != 1:
+                continue
+
+            reshape_input_shape = extract_shape(pre_reshape_input)
+
+            if len(expand_shape) != len(pre_reshape_shape):
+                continue
+
+            # Ensure all dimensions *except* at EXPAND_DIM are identical.
+            if not (
+                expand_shape[:EXPAND_DIM] == pre_reshape_shape[:EXPAND_DIM]
+                and expand_shape[EXPAND_DIM + 1 :]
+                == pre_reshape_shape[EXPAND_DIM + 1 :]
+            ):
+                continue
+
+            # Ensure the expansion dimension is a clean multiple.
+            if expand_shape[EXPAND_DIM] % pre_reshape_shape[EXPAND_DIM] != 0:
+                continue
+
+            expand_ratio = expand_shape[EXPAND_DIM] // pre_reshape_shape[EXPAND_DIM]
+
+            if expand_ratio <= 1:
+                continue
+
+            cat_nodes = []
+
+            for i in range(reshape_input_shape[CAT_DIM]):
+                with graph.inserting_before(expand):
+                    slice_copy_args = (pre_reshape_input, CAT_DIM, i, i + 1, 1)
+                    slice_node = create_node(
+                        graph,
+                        torch.ops.aten.slice.Tensor,
+                        args=slice_copy_args,
+                        origin=expand,
+                    )
+                with graph.inserting_after(slice_node):
+                    cat_args = ([slice_node] * expand_ratio, CAT_DIM)
+                    cat_node = create_node(
+                        graph,
+                        torch.ops.aten.cat.default,
+                        args=cat_args,
+                        origin=expand,
+                    )
+                    cat_nodes.append(cat_node)
+
+            with graph.inserting_after(expand):
+                cat_args = (cat_nodes, CAT_DIM)
+                cat_node = create_node(
+                    graph,
+                    torch.ops.aten.cat.default,
+                    args=cat_args,
+                    origin=expand,
+                )
+                expand.replace_all_uses_with(cat_node)
+
+            modified = True
+            logger.debug(f"{expand.name} is replaced with {cat_node.name} operators")
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)