tico 0.1.0.dev251106__py3-none-any.whl → 0.2.0.dev260122__py3-none-any.whl

tico/serialize/operators/op_circle_shape.py

@@ -0,0 +1,64 @@
+# 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 Dict, List, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import torch._ops
+    import torch.fx
+import torch
+from circle_schema import circle
+
+from tico.serialize.circle_graph import CircleSubgraph
+from tico.serialize.operators.hashable_opcode import OpCode
+from tico.serialize.operators.node_visitor import NodeVisitor, register_node_visitor
+from tico.serialize.operators.utils import create_builtin_operator, get_op_index
+
+
+@register_node_visitor
+class CircleShapeVisitor(NodeVisitor):
+    """
+    Visitor for circle_custom::shape operator.
+
+    This operator extracts the shape of a tensor.
+    It's implemented using Circle's SHAPE operator.
+    """
+
+    target: List[torch._ops.OpOverload] = [
+        torch.ops.circle_custom.shape,
+    ]
+
+    def __init__(self, op_codes: Dict[OpCode, int], graph: CircleSubgraph):
+        super().__init__(op_codes, graph)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+    ) -> circle.Operator.OperatorT:
+        # Args: (input)
+        input_node = node.args[0]
+
+        # SHAPE operator to get the full shape of input
+        op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.SHAPE, self._op_codes
+        )
+
+        inputs = [input_node]
+        outputs = [node]
+        operator = create_builtin_operator(self.graph, op_index, inputs, outputs)
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.ShapeOptions
+        operator.builtinOptions = circle.ShapeOptions.ShapeOptionsT()
+        operator.builtinOptions.outType = circle.TensorType.TensorType.INT32
+
+        return operator

tico/serialize/operators/op_dequantize_per_channel.py

@@ -19,6 +19,7 @@ if TYPE_CHECKING:
     import torch.fx
 import numpy as np
 import torch
+import torch.ao.quantization.fx._decomposed  # register `dequantize_per_channel`
 from circle_schema import circle
 
 from tico.serialize.circle_graph import CircleSubgraph

tico/serialize/operators/op_dequantize_per_tensor.py

@@ -18,6 +18,7 @@ if TYPE_CHECKING:
     import torch._ops
     import torch.fx
 import torch
+import torch.ao.quantization.fx._decomposed  # register `dequantize_per_tensor`
 from circle_schema import circle
 
 from tico.serialize.circle_graph import CircleSubgraph

tico/serialize/operators/op_transpose_conv.py

@@ -66,6 +66,29 @@ class TransposeConvVisitor(NodeVisitor):
         set_transpose_conv_option(operator, stride)
         return operator
 
+    def define_slice_node(
+        self,
+        src_tensor,
+        begin_vals: List[int],
+        size_vals: List[int],
+        dst_tensor,
+    ) -> circle.Operator.OperatorT:
+        slice_op_index = get_op_index(
+            circle.BuiltinOperator.BuiltinOperator.SLICE, self._op_codes
+        )
+
+        # Begin / Size as int32 const tensors
+        begin_arr = circle_legalize_dtype_to(begin_vals, dtype=torch.int32)
+        size_arr = circle_legalize_dtype_to(size_vals, dtype=torch.int32)
+
+        operator = create_builtin_operator(
+            self.graph, slice_op_index, [src_tensor, begin_arr, size_arr], [dst_tensor]
+        )
+        operator.builtinOptionsType = circle.BuiltinOptions.BuiltinOptions.SliceOptions
+        option = circle.SliceOptions.SliceOptionsT()
+        operator.builtinOptions = option
+        return operator
+
     def __init__(self, op_codes: Dict[OpCode, int], graph):
         super().__init__(op_codes, graph)
 
@@ -88,65 +111,58 @@ class TransposeConvVisitor(NodeVisitor):
         assert len(output_shape) == 4, len(output_shape)
         assert len(weight_shape) == 4, len(weight_shape)
 
-        pad_decision = identify_padding(padding, input_shape, output_shape, stride)
+        pad_decision = identify_padding(
+            padding, input_shape, output_shape, stride, is_transpose=True
+        )
 
         conv_input: torch.fx.Node | circle.Tensor.TensorT = input_
-        if pad_decision.explicit_pad_hw is not None:
-            pad_h, pad_w = pad_decision.explicit_pad_hw
-            paddings = torch.tensor(
-                [
-                    [0, 0],
-                    [pad_h, pad_h],
-                    [pad_w, pad_w],
-                    [0, 0],
-                ],
-                dtype=torch.int32,
-            )
-            pad_output_shape: List[int | torch.SymInt] = [
-                input_shape[0],
-                input_shape[1] + pad_h * 2,
-                input_shape[2] + pad_w * 2,
-                input_shape[3],
-            ]
-            pad_output_cshape, pad_output_cshape_signature = to_circle_shape(
-                pad_output_shape
-            )
-            # create padded output tensor
-            input_qparam: Optional[QuantParam] = input_.meta.get(QPARAM_KEY)
-            pad_output = self.graph.add_tensor_from_scratch(
-                prefix=f"{node.name}_input_pad_output",
-                shape=pad_output_cshape,
-                shape_signature=pad_output_cshape_signature,
-                dtype=extract_circle_dtype(input_),
-                qparam=input_qparam,
-                source_node=node,
-            )
-            # CirclePad
-            pad_operator = define_pad_node(
-                self.graph, self._op_codes, [input_, paddings], [pad_output]
-            )
-            self.graph.add_operator(pad_operator)
-            conv_input = pad_output
-
         if bias is None:
             # luci-interpreter can't run no bias conv. Let's add zero vector for bias.
             bias = [0.0] * weight_shape[0]  # type: ignore[assignment]
 
-        # First arguemnt is output shape of tconv.
-        assert output_shape[0] == input_shape[0]
-        assert output_shape[3] == weight_shape[0]
-        tconv_output = circle_legalize_dtype_to(output_shape, dtype=torch.int32)
-
-        tconv_output_tensor = self.graph.add_const_tensor(
-            tconv_output, source_node=node
-        )
+        # Compute pre-crop output shape if we need to apply an explicit crop.
+        if pad_decision.output_crop_hw is not None:
+            pad_h, pad_w = pad_decision.output_crop_hw
+            pre_h = int(output_shape[1]) + 2 * pad_h
+            pre_w = int(output_shape[2]) + 2 * pad_w
+            pre_out_shape = [output_shape[0], pre_h, pre_w, output_shape[3]]  # NHWC
+        else:
+            pre_out_shape = list(output_shape)
+
+        tconv_output = circle_legalize_dtype_to(pre_out_shape, dtype=torch.int32)
+
+        pre_out_cshape, pre_out_csig = to_circle_shape(pre_out_shape)
+        tconv_tmp = node  # type: ignore[assignment]
+        if pad_decision.output_crop_hw is not None:
+            tconv_tmp = self.graph.add_tensor_from_scratch(  # type: ignore[assignment]
+                prefix=f"{node.name}_tconv_out_pre_crop",
+                shape=pre_out_cshape,
+                shape_signature=pre_out_csig,
+                dtype=extract_circle_dtype(node),
+                qparam=node.meta.get(QPARAM_KEY),
+                source_node=node,
+            )
 
-        # TConv2D
         tconv2d_operator = self.define_transpose_conv_node(
-            pad_decision.conv_padding_type,  # 'SAME'(0) or 'VALID'(1)
+            pad_decision.conv_padding_type,
             stride,
-            [tconv_output_tensor, weight, conv_input, bias],
-            [node],
+            [tconv_output, weight, conv_input, bias],
+            [tconv_tmp],
         )
 
+        # If we need an output crop, insert a SLICE to produce the final tensor.
+        if pad_decision.output_crop_hw is not None:
+            self.graph.add_operator(tconv2d_operator)
+            pad_h, pad_w = pad_decision.output_crop_hw
+            begin = [0, pad_h, pad_w, 0]
+            size = [
+                int(output_shape[0]),
+                int(output_shape[1]),
+                int(output_shape[2]),
+                int(output_shape[3]),
+            ]
+
+            slice_op = self.define_slice_node(tconv_tmp, begin, size, node)
+            return slice_op
+
         return tconv2d_operator

tico/utils/convert.py

@@ -25,10 +25,12 @@ 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
+from tico.passes.convert_conv3d_to_conv2d import ConvertConv3dToConv2d
 from tico.passes.convert_expand_to_slice_cat import ConvertExpandToSliceCat
 from tico.passes.convert_layout_op_to_reshape import ConvertLayoutOpToReshape
 from tico.passes.convert_matmul_to_linear import ConvertMatmulToLinear
 from tico.passes.convert_repeat_to_expand_copy import ConvertRepeatToExpandCopy
+from tico.passes.convert_sym_size_to_circle_shape import ConvertSymSizeToCircleShape
 from tico.passes.convert_to_relu6 import ConvertToReLU6
 from tico.passes.decompose_addmm import DecomposeAddmm
 from tico.passes.decompose_batch_norm import DecomposeBatchNorm
@@ -47,6 +49,7 @@ from tico.passes.legalize_causal_mask_value import LegalizeCausalMaskValue
 from tico.passes.legalize_predefined_layout_operators import (
     LegalizePreDefinedLayoutOperators,
 )
+from tico.passes.lower_copy import LowerCopy
 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 passes as LowerToSlicePasses
@@ -97,12 +100,15 @@ def traced_run_decompositions(exported_program: ExportedProgram):
             torch.ops.aten.conv2d.padding,
             torch.ops.aten.conv1d.default,
             torch.ops.aten.conv1d.padding,
+            torch.ops.aten.conv3d.default,
+            torch.ops.aten.conv3d.padding,
             torch.ops.aten.conv_transpose2d.input,
             torch.ops.aten.instance_norm.default,
             torch.ops.aten._safe_softmax.default,
             torch.ops.aten.relu6.default,  # Do not decompose to hardtanh
             torch.ops.aten.linear.default,
             torch.ops.aten.upsample_nearest2d.vec,
+            torch.ops.aten.rms_norm.default,
         )
         ep = ep.run_decompositions(_preserve_ops=_preserve_ops)
 
@@ -115,6 +121,8 @@ def traced_run_decompositions(exported_program: ExportedProgram):
             torch.ops.aten.conv2d.padding,
             torch.ops.aten.conv1d.default,
             torch.ops.aten.conv1d.padding,
+            torch.ops.aten.conv3d.default,
+            torch.ops.aten.conv3d.padding,
             torch.ops.aten.conv_transpose2d.input,
             torch.ops.aten.instance_norm.default,
             torch.ops.aten._safe_softmax.default,
@@ -138,6 +146,7 @@ def traced_run_decompositions(exported_program: ExportedProgram):
         or torch.__version__.startswith("2.8")
         or torch.__version__.startswith("2.9")
         or torch.__version__.startswith("2.10")
+        or torch.__version__.startswith("2.11")
     ):
         return run_decompositions(exported_program)
     else:
@@ -224,6 +233,8 @@ def convert_exported_module_to_circle(
             FillMetaVal(),
             ExtractDtypeKwargsPass(),
             RemoveNop(),
+            LowerCopy(),
+            ConvertSymSizeToCircleShape(),
             ConvertLayoutOpToReshape(),
             RestoreLinear(),
             ConvertToReLU6(),
@@ -258,6 +269,7 @@ def convert_exported_module_to_circle(
             LegalizePreDefinedLayoutOperators(),
             LowerPow2ToMul(),
             ConvertConv1dToConv2d(),
+            ConvertConv3dToConv2d(),
             *LowerToSlicePasses(),
             FuseLeadingUnsqueezeReshape(),
             CastClampMixedTypeArgs(),
@@ -316,7 +328,10 @@ def convert(
             mod, args, kwargs, dynamic_shapes=dynamic_shapes, strict=strict
         )
 
-    circle_binary = convert_exported_module_to_circle(exported_program, config=config)
+    with SuppressWarning(FutureWarning, ".*LeafSpec*"):
+        circle_binary = convert_exported_module_to_circle(
+            exported_program, config=config
+        )
 
     return CircleModel(circle_binary)
 

tico/utils/padding.py

@@ -35,6 +35,7 @@ class ConvPaddingInfo(NamedTuple):
 
     conv_padding_type: ConvPadding
     explicit_pad_hw: Optional[Tuple[int, int]]  # None -> no extra Pad() op needed
+    output_crop_hw: Optional[Tuple[int, int]]
 
 
 def identify_padding(
@@ -42,6 +43,7 @@ def identify_padding(
     input_shape: Sequence[int | torch.SymInt] | torch.Size,
     output_shape: Sequence[int | torch.SymInt] | torch.Size,
     stride: Sequence[int],
+    is_transpose: bool = False,
 ) -> ConvPaddingInfo:
     """
     Normalizes all PyTorch `padding` variants to a single decision.
@@ -61,9 +63,9 @@ def identify_padding(
     if isinstance(padding, str):
         pad = padding.lower()
         if pad == "valid":
-            return ConvPaddingInfo(ConvPadding.VALID, None)
+            return ConvPaddingInfo(ConvPadding.VALID, None, None)
         if pad == "same":
-            return ConvPaddingInfo(ConvPadding.SAME, None)
+            return ConvPaddingInfo(ConvPadding.SAME, None, None)
         raise InvalidArgumentError(f"Unknown padding string: {padding}")
 
     # ─── 2. List / tuple form ─────────────────────────────────────────────
@@ -73,15 +75,21 @@ def identify_padding(
         )
 
     pad_h, pad_w = padding
+
+    if is_transpose:
+        if pad_h == 0 and pad_w == 0:
+            return ConvPaddingInfo(ConvPadding.VALID, None, None)
+        return ConvPaddingInfo(ConvPadding.VALID, None, (pad_h, pad_w))
+
     # [0, 0]  → VALID
     if pad_h == 0 and pad_w == 0:
-        return ConvPaddingInfo(ConvPadding.VALID, None)
+        return ConvPaddingInfo(ConvPadding.VALID, None, None)
 
     # SAME heuristic: output H/W already match input when stride is 1
     hw_in = tuple(input_shape[1:3])
     hw_out = tuple(output_shape[1:3])
     if hw_in == hw_out and stride == [1, 1]:
-        return ConvPaddingInfo(ConvPadding.SAME, None)
+        return ConvPaddingInfo(ConvPadding.SAME, None, None)
 
     # Anything else = explicit symmetric padding
-    return ConvPaddingInfo(ConvPadding.VALID, (pad_h, pad_w))
+    return ConvPaddingInfo(ConvPadding.VALID, (pad_h, pad_w), None)

tico/utils/record_input.py

@@ -49,13 +49,13 @@ class RecordingInput:
     def __init__(
         self,
         module: nn.Module,
-        condition: Callable[[dict], bool] = lambda args_dict: True,
+        condition: Optional[Callable[[dict], bool]] = None,
         *,
         input_to_remove: Optional[List[str]] = [],
     ):
         self.module = module
         self.forward_org = module.forward
-        self.condition = condition
+        self.condition = condition or (lambda args_dict: True)
         self.input_to_remove = input_to_remove
         self.sig = inspect.signature(self.forward_org)
 

tico/utils/register_custom_op.py

@@ -727,6 +727,67 @@ def CircleRMSNorm():
         return hidden_states.new_empty(hidden_states.size())
 
 
+def CircleAttention():
+    @custom_op("circle_custom::attention", mutates_args=())
+    def attention(
+        hidden_states: torch.Tensor,
+        wq: torch.Tensor,
+        wk: torch.Tensor,
+        wv: torch.Tensor,
+        wo: torch.Tensor,
+        position_cos: torch.Tensor,
+        position_sin: torch.Tensor,
+        attention_mask: torch.Tensor,
+        past_key: torch.Tensor,
+        past_value: torch.Tensor,
+        cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        return None
+
+    @register_fake("circle_custom::attention")
+    def _(
+        hidden_states: torch.Tensor,
+        wq: torch.Tensor,
+        wk: torch.Tensor,
+        wv: torch.Tensor,
+        wo: torch.Tensor,
+        position_cos: torch.Tensor,
+        position_sin: torch.Tensor,
+        attention_mask: torch.Tensor,
+        past_key: torch.Tensor,
+        past_value: torch.Tensor,
+        cache_position: torch.Tensor,
+    ) -> torch.Tensor:
+        return hidden_states
+
+
+def CircleShape():
+    """
+    Custom operator to extract the shape of a tensor.
+    This is similar to TensorFlow's shape operator and is used to preserve
+    dynamic shape information in the Circle model.
+
+    Args:
+        input_: Input tensor
+
+    Returns:
+        A 1D tensor containing the shape of the input tensor
+    """
+
+    @custom_op("circle_custom::shape", mutates_args=())
+    def shape(input_: torch.Tensor) -> torch.Tensor:
+        # Return the shape of the input tensor as a 1D tensor
+        shape_val = list(input_.size())
+        return torch.tensor(shape_val, dtype=torch.int32)
+
+    @register_fake("circle_custom::shape")
+    def _(input_: torch.Tensor) -> torch.Tensor:
+        # Return a 1D tensor with symbolic shape
+        # The actual value will be determined at runtime
+        rank = len(input_.size())
+        return torch.empty([rank], dtype=torch.int32)
+
+
 # Add custom ops to the torch namespace
 def RegisterOps():
     CircleResizeNearestNeighbor()
@@ -740,3 +801,5 @@ def RegisterOps():
     CircleInstanceNorm()
     CircleQuantizeMX()
     CircleRMSNorm()
+    CircleAttention()
+    CircleShape()

tico/utils/validate_args_kwargs.py

@@ -171,6 +171,26 @@ class CatArgs:
     dim: int = 0
 
 
+@enforce_type
+@dataclass
+class CircleAttentionArgs:
+    """
+    For circle.BuiltinOperator.BuiltinOperator.ATTENTION
+    """
+
+    hidden_states: torch.fx.Node
+    wq: torch.fx.Node
+    wk: torch.fx.Node
+    wv: torch.fx.Node
+    wo: torch.fx.Node
+    position_cos: torch.fx.Node
+    position_sin: torch.fx.Node
+    attention_mask: torch.fx.Node
+    past_key: torch.fx.Node
+    past_value: torch.fx.Node
+    cache_position: torch.fx.Node
+
+
 @enforce_type
 @dataclass
 class CircleRMSNormArgs:
@@ -299,6 +319,27 @@ class Conv1DArgs:
         assert len(self.dilation) == 1, len(self.dilation)
 
 
+@enforce_type
+@dataclass
+class Conv3DArgs:
+    """
+    conv3d(Tensor input, Tensor weight, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, SymInt[3] dilation=1, SymInt groups=1) -> Tensor
+    conv3d.padding(Tensor input, Tensor weight, Tensor? bias=None, SymInt[3] stride=1, str padding="valid", SymInt[3] dilation=1, SymInt groups=1) -> Tensor
+    """
+
+    input: torch.fx.Node
+    weight: torch.fx.Node
+    bias: Union[torch.fx.Node, None] = None
+    stride: List[int] = field(default_factory=lambda: [1, 1, 1])
+    padding: Union[List[int], str] = field(default_factory=lambda: [0, 0, 0])
+    dilation: List[int] = field(default_factory=lambda: [1, 1, 1])
+    groups: int = 1
+
+    def __post_init__(self):
+        assert len(self.stride) == 3, len(self.stride)
+        assert len(self.dilation) == 3, len(self.dilation)
+
+
 @enforce_type
 @dataclass
 class CopyArgs:
@@ -930,7 +971,7 @@ class RepeatArgs:
     """
 
     input: torch.fx.Node
-    repeats: List[int]
+    repeats: List[Union[int, torch.SymInt, torch.fx.Node]]
 
 
 @enforce_type
@@ -938,10 +979,14 @@ class RepeatArgs:
 class ReshapeArgs:
     """
     reshape(Tensor(a) self, SymInt[] shape) -> Tensor(a)
+
+    Note: After PrepareReshapeDynamicShape pass, shape can be either:
+    - A list of int/SymInt/Node (original or static)
+    - A single Node (dynamic shape tensor prepared by the pass)
     """
 
     input: torch.fx.Node
-    shape: List[int]
+    shape: Union[List[Union[int, torch.SymInt, torch.fx.Node]], torch.fx.Node]
 
 
 @enforce_type
@@ -1077,7 +1122,7 @@ class SplitWithSizesArgs:
     """
 
     input: torch.fx.Node
-    split_sizes: List[int]
+    split_sizes: List[Union[int, torch.SymInt, torch.fx.Node]]
     dim: int = 0
 
 
@@ -1218,7 +1263,7 @@ class ViewArgs:
     """
 
     input: torch.fx.Node
-    size: List[int]
+    size: List[Union[int, torch.SymInt, torch.fx.Node]]
 
 
 @enforce_type