tico 0.1.0.dev250904__py3-none-any.whl → 0.1.0.dev251109__py3-none-any.whl

tico/__init__.py

@@ -29,7 +29,7 @@ __all__ = [
 ]
 
 # THIS LINE IS AUTOMATICALLY GENERATED BY setup.py
-__version__ = "0.1.0.dev250904"
+__version__ = "0.1.0.dev251109"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

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_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,

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)

tico/passes/convert_matmul_to_linear.py

@@ -0,0 +1,312 @@
+# 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 List, Optional, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch.fx
+import torch
+from torch._export.utils import is_buffer, is_lifted_tensor_constant, is_param
+from torch.export import ExportedProgram
+
+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.validate_args_kwargs import BmmArgs, MatmulArgs
+
+
+class Converter:  # type: ignore[empty-body]
+    def __init__(self):
+        super().__init__()
+
+    def match(self, exported_program, node) -> bool:  # type: ignore[empty-body]
+        return False
+
+    def convert(self, exported_program, node) -> torch.fx.Node:  # type: ignore[empty-body]
+        pass
+
+
+class MatmulToLinearConverter(Converter):
+    def __init__(self):
+        super().__init__()
+
+    def convert(self, exported_program, node) -> torch.fx.Node:
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+
+        mm_args = MatmulArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+
+        lhs = mm_args.input
+        rhs = mm_args.other
+
+        with graph.inserting_before(node):
+            transpose_node = create_node(
+                graph,
+                torch.ops.aten.permute.default,
+                args=(rhs, [1, 0]),
+            )
+            linear_node = create_node(
+                graph,
+                torch.ops.aten.linear.default,
+                args=(lhs, transpose_node),
+            )
+            node.replace_all_uses_with(linear_node, propagate_meta=True)
+
+        return linear_node
+
+
+class RhsConstMatmulToLinearConverter(MatmulToLinearConverter):
+    def __init__(self):
+        super().__init__()
+
+    def match(self, exported_program, node) -> bool:
+        if not node.target == torch.ops.aten.mm.default:
+            return False
+
+        mm_args = MatmulArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+
+        rhs = mm_args.other
+        if isinstance(rhs, torch.fx.Node):
+            if is_lifted_tensor_constant(exported_program, rhs):
+                return True
+            elif is_param(exported_program, rhs):
+                return True
+            elif is_buffer(exported_program, rhs):
+                return True
+            else:
+                return False
+        return False
+
+    def convert(self, exported_program, node) -> torch.fx.Node:
+        return super().convert(exported_program, node)
+
+
+class LhsConstMatmulToLinearConverter(MatmulToLinearConverter):
+    def __init__(self):
+        super().__init__()
+
+    def match(self, exported_program, node) -> bool:
+        if not node.target == torch.ops.aten.mm.default:
+            return False
+
+        mm_args = MatmulArgs(*node.args, **node.kwargs)
+        lhs = mm_args.input
+        if isinstance(lhs, torch.fx.Node):
+            if is_lifted_tensor_constant(exported_program, lhs):
+                return True
+            elif is_param(exported_program, lhs):
+                return True
+            elif is_buffer(exported_program, lhs):
+                return True
+        return False
+
+    def convert(self, exported_program, node) -> torch.fx.Node:
+        return super().convert(exported_program, node)
+
+
+class SingleBatchLhsConstBmmToLinearConverter(Converter):
+    """
+    Convert `single-batched & lhs-const BatchMatMul` to `linear` operation.
+
+    [1] exchange lhs and rhs
+    [2] transpose rhs
+    [3] transpose output
+
+    **Before**
+
+    lhs[1,a,b](const)   rhs[1,b,c]
+    |                   |
+    |                   |
+    ---------bmm---------
+              |
+            output[1,a,c]
+
+
+    **After**
+
+    rhs[1,b,c]
+    |
+    tr                  lhs'[a,b](const-folded)
+    |[1,c,b]            |
+    |                   |
+    ---------fc--------
+             |[1,c,a]
+             tr
+             |
+            output[1,a,c]
+
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def match(self, exported_program, node) -> bool:
+        if not node.target == torch.ops.aten.bmm.default:
+            return False
+
+        bmm_args = BmmArgs(*node.args, **node.kwargs)
+        lhs = bmm_args.input
+        rhs = bmm_args.mat2
+
+        # [1] Single-batch
+        lhs_shape = extract_shape(lhs)
+        rhs_shape = extract_shape(rhs)
+
+        assert len(lhs_shape) == len(
+            rhs_shape
+        ), f"Bmm input's ranks must be the same but got {lhs_shape} and {rhs_shape}"
+
+        if not (lhs_shape[0] == rhs_shape[0] == 1):
+            return False
+
+        # [2] Lhs is constant
+        if not isinstance(lhs, torch.fx.Node):
+            return False
+        if not (
+            is_lifted_tensor_constant(exported_program, lhs)
+            or is_param(exported_program, lhs)
+            or is_buffer(exported_program, lhs)
+        ):
+            return False
+
+        return True
+
+    def convert(self, exported_program, node) -> torch.fx.Node:
+        graph_module = exported_program.graph_module
+        graph = graph_module.graph
+
+        bmm_args = BmmArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+
+        lhs = bmm_args.input  # const
+        rhs = bmm_args.mat2  # non-const
+        lhs_shape = extract_shape(lhs)
+        rhs_shape = extract_shape(rhs)
+        assert rhs_shape[0] == 1
+        assert lhs_shape[0] == 1
+
+        with graph.inserting_before(node):
+            rhs_tr = create_node(
+                graph,
+                torch.ops.aten.permute.default,
+                args=(rhs, [0, 2, 1]),
+            )
+            lhs_reshape = create_node(
+                graph,
+                torch.ops.aten.view.default,
+                args=(lhs, list(lhs_shape[1:])),
+            )
+
+            linear_node = create_node(
+                graph,
+                torch.ops.aten.linear.default,
+                args=(rhs_tr, lhs_reshape),
+            )
+
+            tr_linear_node = create_node(
+                graph,
+                torch.ops.aten.permute.default,
+                args=(linear_node, [0, 2, 1]),
+            )
+
+            node.replace_all_uses_with(tr_linear_node, propagate_meta=False)
+
+        return tr_linear_node
+
+
+@trace_graph_diff_on_pass
+class ConvertMatmulToLinear(PassBase):
+    """
+    This pass converts matmul(partially includes single-batch bmm) to linear selectively
+
+    How to select between `matmul` and `linear`?
+
+    * Linear has better quantization accuracy (NPU backend)
+        Due to ONE compiler's quantization policy;
+        FullyConnected(=Linear) uses per-channel quantization for weight and per-tensor for input.
+        BatchMatmul(=matmul) uses per-tensor quantization for both rhs and lhs.
+
+    * Matmul to Linear requires Transpose, which may harm latency
+        When RHS is constant, addtional transpose can be folded.
+
+    [RHS non-const case]
+    Constant folding cannot be performed.
+
+    lhs         rhs (non-const)
+    |           |
+    |           transpose
+    |           |
+     -- linear --
+         |
+         out
+
+    [RHS const case]
+    Constant folding can be performed to
+
+    lhs         rhs (const)         lh          rhs (folded const)
+    |           |                   |           |
+    |           transpose           |           |
+    |           |                   |           |
+     -- linear --         -->        -- linear --
+         |                                |
+         out                              out
+
+
+    enable_lhs_const: If true, convert matmul where LHS is constant tensor. Default is False.
+    enable_rhs_const: If true, convert matmul where RHS is constant tensor. Default is True.
+    """
+
+    def __init__(
+        self,
+        enable_lhs_const: Optional[bool] = False,
+        enable_rhs_const: Optional[bool] = True,
+        enable_single_batch_lhs_const_bmm: Optional[bool] = False,
+    ):
+        super().__init__()
+        self.converters: List[Converter] = []
+        if enable_lhs_const:
+            self.converters.append(LhsConstMatmulToLinearConverter())
+        if enable_rhs_const:
+            self.converters.append(RhsConstMatmulToLinearConverter())
+        if enable_single_batch_lhs_const_bmm:
+            self.converters.append(SingleBatchLhsConstBmmToLinearConverter())
+
+    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
+
+            for converter in self.converters:
+                if not converter.match(exported_program, node):
+                    continue
+
+                new_node = converter.convert(exported_program, node)
+                modified = True
+                logger.debug(
+                    f"{node.name} is replaced with {new_node.name} operator (permute + linear)"
+                )
+                continue
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        return PassResult(modified)

tico/passes/convert_to_relu6.py

@@ -172,7 +172,7 @@ class ConvertToReLU6(PassBase):
                 converter.convert(exported_program, node)
                 modified = True
                 logger.debug(f"{node.name} is replaced with ReLU6 operator")
-                break
+                continue
 
         graph.eliminate_dead_code()
         graph.lint()

tico/passes/decompose_fake_quantize_tensor_qparams.py

@@ -245,9 +245,10 @@ class DecomposeFakeQuantizeTensorQParams(PassBase):
                     # mask_user(output).args == (dequantize_per_tensor.tensor, mask)
                     if mask:
                         assert len(mask) == 1
-                        mask_user = list(mask[0].users.keys())[0]
-                        assert len(mask_user.args) == 1
-                        mask_user.args = ((mask_user.args[0][0],),)
+                        if len(mask[0].users) > 0:
+                            mask_user = list(mask[0].users.keys())[0]
+                            assert len(mask_user.args) == 1
+                            mask_user.args = ((mask_user.args[0][0],),)
                 modified = True
             if (
                 node.target

tico/passes/ops.py

@@ -69,7 +69,6 @@ class AtenOps:
             torch.ops.aten.unsqueeze_copy.default,
         ]
         self.view = [
-            torch.ops.aten.view,
             torch.ops.aten.view.default,
             torch.ops.aten.view_copy.default,
         ]

tico/passes/remove_redundant_expand.py

@@ -46,7 +46,9 @@ class RemoveRedundantExpand(PassBase):
             input, size = args.input, args.size
 
             input_shape = extract_shape(input)
-            if list(input_shape) != size:
+            output_shape = extract_shape(node)
+
+            if input_shape != output_shape:
                 continue
 
             node.replace_all_uses_with(input, propagate_meta=False)

tico/quantization/__init__.py

@@ -0,0 +1,6 @@
+from tico.quantization.public_interface import convert, prepare
+
+__all__ = [
+    "convert",
+    "prepare",
+]