tico 0.1.0.dev250616__py3-none-any.whl → 0.1.0.dev250618__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.dev250616"
+__version__ = "0.1.0.dev250618"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py

@@ -24,6 +24,7 @@ from torch.export import ExportedProgram
 from tico.serialize.quant_param import QPARAM_KEY, QuantParam
 from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
+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 quant_min_max, set_new_meta_val
@@ -145,9 +146,11 @@ class InsertQuantizeOnDtypeMismatch(PassBase):
 
             with graph.inserting_before(node):
                 q_args = (inp, scale, zerop, min_, max_, dtype)
-                quantize = graph.call_function(
+                quantize = create_node(
+                    graph,
                     torch.ops.quantized_decomposed.quantize_per_tensor.default,
                     args=q_args,
+                    origin=node,
                 )
                 quantize.meta[QPARAM_KEY] = copy.deepcopy(qparam)
                 set_new_meta_val(quantize)
@@ -166,7 +169,8 @@ class InsertQuantizeOnDtypeMismatch(PassBase):
             dtype = getattr(torch, qparam.dtype)
             with graph.inserting_after(node):
                 q_args = (node, scale, zerop, min_, max_, dtype)
-                quantize = graph.call_function(
+                quantize = create_node(
+                    graph,
                     torch.ops.quantized_decomposed.quantize_per_tensor.default,
                     args=q_args,
                 )

tico/experimental/quantization/passes/quantize_bias.py

@@ -0,0 +1,123 @@
+# 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 copy
+
+import torch
+from torch.export import ExportedProgram
+
+from tico.serialize.quant_param import QPARAM_KEY, QuantParam, to_qparam_dtype
+from tico.utils import logging
+from tico.utils.graph import add_placeholder, get_torch_param_value, is_torch_param
+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 LinearArgs
+
+
+@trace_graph_diff_on_pass
+class QuantizeBias(PassBase):
+    """
+    Quantize bias.
+
+    This pass identifies fp32 biases, quantizes them using scales of input and weights.
+
+    This pass assumes that if bias is fp32, input and weights must have been quantized.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def call(self, exported_program: ExportedProgram) -> PassResult:
+        logger = logging.getLogger(__name__)
+
+        graph_module = exported_program.graph_module
+        graph: torch.fx.Graph = graph_module.graph
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target == torch.ops.aten.linear.default:
+                lin_args = LinearArgs(*node.args, **node.kwargs)
+                inp = lin_args.input
+                weights = lin_args.weight
+                bias = lin_args.bias
+
+                if bias is None:
+                    continue
+
+                # Only support bias is Parameter
+                # TODO Is it possible that bias is not Parameter?
+                if not is_torch_param(bias, exported_program):
+                    continue
+
+                bias_val: torch.Tensor = get_torch_param_value(bias, exported_program)
+                if bias_val.dtype != torch.float32:
+                    continue
+
+                if QPARAM_KEY not in inp.meta:
+                    continue
+
+                if QPARAM_KEY not in weights.meta:
+                    continue
+
+                quant_dtype = None
+                if inp.meta[QPARAM_KEY].dtype == "int16":
+                    quant_dtype = torch.int64
+                elif inp.meta[QPARAM_KEY].dtype == "uint8":
+                    quant_dtype = torch.int32
+                else:
+                    continue
+
+                type_info = torch.iinfo(quant_dtype)
+
+                assert quant_dtype is not None
+
+                i_scale = inp.meta[QPARAM_KEY].scale
+                w_scale = weights.meta[QPARAM_KEY].scale
+
+                assert i_scale is not None
+                assert w_scale is not None
+                assert len(i_scale) == 1
+                assert len(w_scale) == bias_val.shape[0]
+
+                bias_scale = torch.tensor(i_scale) * torch.tensor(w_scale)
+                q_bias = torch.round(bias_val / bias_scale)
+                q_bias = torch.clamp(q_bias, min=type_info.min, max=type_info.max)
+                q_bias = q_bias.to(quant_dtype)
+
+                q_bias_node = add_placeholder(exported_program, q_bias, bias.name)
+
+                qparam = QuantParam()
+                qparam.scale = bias_scale.tolist()
+                assert qparam.scale is not None
+                qparam.zero_point = [0] * len(qparam.scale)
+                qparam.dtype = to_qparam_dtype(quant_dtype)
+                qparam.quantized_dimension = 0
+                q_bias_node.meta[QPARAM_KEY] = qparam
+
+                node.update_arg(2, q_bias_node)
+
+                logger.debug(f"Bias ({bias.name}) is quantized to {q_bias_node.name}.")
+
+            # TODO Support more ops.
+
+        graph.eliminate_dead_code()
+        graph.lint()
+        graph_module.recompile()
+
+        # Run only once.
+        return PassResult(False)

tico/experimental/quantization/passes/remove_weight_dequant_op.py

@@ -145,6 +145,9 @@ class RemoveWeightDequantOp(PassBase):
             if isinstance(dq_args, DequantizePerChannelArgs):
                 scales = get_constant(exported_program, dq_args.scales)
                 zero_ps = get_constant(exported_program, dq_args.zero_points)
+
+                # Sometimes users can give fp32 zero point. Let's update dtype here.
+                zero_ps = zero_ps.to(torch.int64)
                 quant_param.scale = scales.tolist()
                 quant_param.zero_point = zero_ps.tolist()
                 assert quant_param.zero_point is not None  # To avoid mypy error

tico/passes/cast_aten_where_arg_type.py

@@ -21,6 +21,7 @@ from torch.export import ExportedProgram
 
 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_const_diff_on_pass,
@@ -158,10 +159,12 @@ class CastATenWhereArgType(PassBase):
                         f"{to_cast.name}({buf_data.dtype}) data range is out of {dtype_to_cast} range"
                     )
             with graph_module.graph.inserting_after(to_cast):
-                cast = graph_module.graph.call_function(
+                cast = create_node(
+                    graph,
                     torch.ops.aten._to_copy.default,
                     args=(to_cast,),
                     kwargs={"dtype": dtype_to_cast},
+                    origin=to_cast,
                 )
             # set new meta["val"] in advance because we will use it below for checking if type promotion is valid.
             set_new_meta_val(cast)

tico/passes/cast_mixed_type_args.py

@@ -26,6 +26,7 @@ from torch.export import ExportedProgram
 
 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
@@ -126,10 +127,12 @@ class CastMixedTypeArgs(PassBase):
 
             if isinstance(arg_to_promote, torch.fx.Node):
                 with graph.inserting_after(arg_to_promote):
-                    to_copy = graph.call_function(
+                    to_copy = create_node(
+                        graph,
                         torch.ops.aten._to_copy.default,
                         (arg_to_promote,),
                         {"dtype": type_to_promote},
+                        origin=arg_to_promote,
                     )
                     # set new meta["val"] in advance because we will use it below for checking if type promotion is valid.
                     set_new_meta_val(to_copy)

tico/passes/convert_conv1d_to_conv2d.py

@@ -22,6 +22,7 @@ from torch.export import ExportedProgram
 from tico.serialize.circle_graph import extract_shape
 from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
+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
@@ -89,15 +90,19 @@ class ConvertConv1dToConv2d(PassBase):
             )
 
         with graph.inserting_after(input):
-            input_unsqueeze = graph_module.graph.call_function(
+            input_unsqueeze = create_node(
+                graph,
                 torch.ops.aten.unsqueeze.default,
                 args=(input, 3),
+                origin=input,
             )
 
         with graph.inserting_after(weight):
-            weight_unsqueeze = graph_module.graph.call_function(
+            weight_unsqueeze = create_node(
+                graph,
                 torch.ops.aten.unsqueeze.default,
                 args=(weight, 3),
+                origin=weight,
             )
 
         with graph.inserting_before(node):
@@ -106,7 +111,8 @@ class ConvertConv1dToConv2d(PassBase):
             elif isinstance(padding, str):
                 conv2d_op = torch.ops.aten.conv2d.padding
 
-            conv2d = graph_module.graph.call_function(
+            conv2d = create_node(
+                graph,
                 conv2d_op,
                 args=(
                     input_unsqueeze,
@@ -118,9 +124,11 @@ class ConvertConv1dToConv2d(PassBase):
                     groups,
                 ),
                 kwargs=node.kwargs,
+                origin=node,
             )
 
-            conv_out_squeeze = graph_module.graph.call_function(
+            conv_out_squeeze = create_node(
+                graph,
                 torch.ops.aten.squeeze.dims,
                 args=(conv2d, [3]),
             )

tico/passes/convert_layout_op_to_reshape.py

@@ -22,6 +22,7 @@ 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.validate_args_kwargs import SqueezeArgs, UnSqueezeArgs, ViewArgs
@@ -48,11 +49,11 @@ class ConvertLayoutOpToReshape(PassBase):
             out_shape = list(extract_shape(node))
 
             with graph.inserting_after(node):
-                reshape_node = graph.call_function(
+                reshape_node = create_node(
+                    graph,
                     torch.ops.aten.reshape.default,
                     args=(input, out_shape),
                 )
-
             node.replace_all_uses_with(reshape_node, propagate_meta=True)
 
             logger.debug(f"{node.name} is replaced with {reshape_node.name}")

tico/passes/convert_repeat_to_expand_copy.py

@@ -20,6 +20,7 @@ import torch
 from torch.export import ExportedProgram
 
 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
@@ -71,8 +72,10 @@ class ConvertRepeatToExpandCopy(PassBase):
             expand_copy_args = (tensor, size)
 
             with graph.inserting_after(node):
-                expand_copy_node = graph.call_function(
-                    torch.ops.aten.expand_copy.default, args=expand_copy_args
+                expand_copy_node = create_node(
+                    graph,
+                    torch.ops.aten.expand_copy.default,
+                    args=expand_copy_args,
                 )
                 node.replace_all_uses_with(expand_copy_node, propagate_meta=True)
 

tico/passes/convert_to_relu6.py

@@ -20,6 +20,7 @@ import torch
 from torch.export import ExportedProgram
 
 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 ClampArgs, HardTanhArgs
@@ -58,7 +59,7 @@ class ConvertHardTanhToReLU6(Converter):
         input = args.input
 
         with graph.inserting_after(node):
-            relu_node = graph.call_function(torch.ops.aten.relu6.default, args=(input,))
+            relu_node = create_node(graph, torch.ops.aten.relu6.default, args=(input,))
             node.replace_all_uses_with(relu_node, propagate_meta=True)
 
 
@@ -84,7 +85,7 @@ class ConvertClampToReLU6(Converter):
         input = args.input
 
         with graph.inserting_after(node):
-            relu_node = graph.call_function(torch.ops.aten.relu6.default, args=(input,))
+            relu_node = create_node(graph, torch.ops.aten.relu6.default, args=(input,))
             node.replace_all_uses_with(relu_node, propagate_meta=True)
 
 
@@ -140,7 +141,7 @@ class ConvertDoubleClampsToReLU6(Converter):
         input = prev_args.input
 
         with graph.inserting_after(node):
-            relu_node = graph.call_function(torch.ops.aten.relu6.default, args=(input,))
+            relu_node = create_node(graph, torch.ops.aten.relu6.default, args=(input,))
             node.replace_all_uses_with(relu_node, propagate_meta=True)
 
 

tico/passes/decompose_addmm.py

@@ -21,7 +21,7 @@ from torch.export import ExportedProgram
 
 from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
-from tico.utils.graph import add_placeholder
+from tico.utils.graph import add_placeholder, 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
@@ -78,7 +78,9 @@ class DecomposeAddmm(PassBase):
 
             with graph.inserting_before(node):
                 # out = beta * input + alpha * (mat1 @ mat2)
-                matmul = graph.call_function(torch.ops.aten.mm.default, (mat1, mat2))
+                matmul = create_node(
+                    graph, torch.ops.aten.mm.default, (mat1, mat2), origin=node
+                )
                 set_new_meta_val(matmul)
 
                 if beta == 1:
@@ -90,7 +92,9 @@ class DecomposeAddmm(PassBase):
                         f"{node.name}_beta_zeros",
                     )
                 else:
-                    bias = graph.call_function(torch.ops.aten.mul.Tensor, (input, beta))
+                    bias = create_node(
+                        graph, torch.ops.aten.mul.Tensor, (input, beta), origin=node
+                    )
 
                 if alpha == 1:
                     scaled_matmul: torch.fx.Node | torch.Tensor = matmul
@@ -101,12 +105,12 @@ class DecomposeAddmm(PassBase):
                         f"{node.name}_alpha_zeros",
                     )
                 else:
-                    scaled_matmul = graph.call_function(
-                        torch.ops.aten.mul.Tensor, (matmul, alpha)
+                    scaled_matmul = create_node(
+                        graph, torch.ops.aten.mul.Tensor, (matmul, alpha), origin=node
                     )
 
-                result = graph.call_function(
-                    torch.ops.aten.add.Tensor, (bias, scaled_matmul)
+                result = create_node(
+                    graph, torch.ops.aten.add.Tensor, (bias, scaled_matmul)
                 )
 
             node.replace_all_uses_with(result, propagate_meta=True)

tico/passes/decompose_batch_norm.py

@@ -24,6 +24,7 @@ from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 from tico.utils.graph import (
     add_placeholder,
+    create_node,
     get_first_user_input,
     get_torch_buffer_value,
     get_torch_param_value,
@@ -32,16 +33,10 @@ from tico.utils.graph import (
 )
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
-from tico.utils.utils import fill_meta_val, is_target_node
+from tico.utils.utils import is_target_node
 from tico.utils.validate_args_kwargs import NativeBatchNormLegitNoTrainingArgs
 
 
-def insert_node(graph: torch.fx.Graph, operation, args):
-    new_node = graph.call_function(operation, args)
-
-    return new_node
-
-
 @trace_graph_diff_on_pass
 class DecomposeBatchNorm(PassBase):
     """
@@ -173,19 +168,20 @@ class DecomposeBatchNorm(PassBase):
                 )
 
             with gm.graph.inserting_before(node):
-                mul = graph.call_function(
+                mul = create_node(
+                    graph,
                     torch.ops.aten.mul.Tensor,
                     args=(input_, mul_const_node),
+                    origin=node,
                 )
-                add = graph.call_function(
+                add = create_node(
+                    graph,
                     torch.ops.aten.add.Tensor,
                     args=(mul, add_const_node),
                 )
-                # Not set meta for propagating replacing get_item's meta.
             get_item, *_ = node.users.keys()
             get_item.replace_all_uses_with(add, propagate_meta=True)
 
-            fill_meta_val(exported_program)
             logger.debug(f"{node.name} is decomposed to {mul.name} and {add.name}")
             modified = True
 

tico/passes/decompose_fake_quantize.py

@@ -23,6 +23,7 @@ from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib
 from torch.export import ExportedProgram
 
 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 FakeQuantizePerChannelArgs
@@ -69,6 +70,7 @@ class DecomposeFakeQuantize(PassBase):
         modified = False
 
         gm = exported_program.graph_module
+        g = gm.graph
         qd = torch.ops.quantized_decomposed  # type: ignore[return]
         for node in gm.graph.nodes:
             if node.op != "call_function":
@@ -83,17 +85,19 @@ class DecomposeFakeQuantize(PassBase):
                     **{"dtype": get_quant_type(quant_min, quant_max)},
                 }
                 with gm.graph.inserting_before(node):
-                    quant = gm.graph.call_function(
+                    quant = create_node(
+                        g,
                         qd.quantize_per_tensor.default,
                         args=node.args,
                         kwargs=quant_kwargs,
+                        origin=node,
                     )
-                    dequnt = gm.graph.call_function(
+                    dequnt = create_node(
+                        g,
                         qd.dequantize_per_tensor.default,
                         args=(quant, *quant.args[1:]),
                         kwargs=quant.kwargs,
                     )
-                    # Not set meta for propagating replacing node's meta.
                     node.replace_all_uses_with(dequnt, propagate_meta=True)
                 modified = True
 
@@ -107,17 +111,19 @@ class DecomposeFakeQuantize(PassBase):
                     **{"dtype": get_quant_type(quant_min, quant_max)},
                 }
                 with gm.graph.inserting_before(node):
-                    quant = gm.graph.call_function(
+                    quant = create_node(
+                        g,
                         qd.quantize_per_channel.default,
                         args=node.args,
                         kwargs=quant_kwargs,
+                        origin=node,
                     )
-                    dequnt = gm.graph.call_function(
+                    dequnt = create_node(
+                        g,
                         qd.dequantize_per_channel.default,
                         args=(quant, *quant.args[1:]),
                         kwargs=quant.kwargs,
                     )
-                    # Not set meta for propagating replacing node's meta.
                     node.replace_all_uses_with(dequnt, propagate_meta=True)
                 modified = True
 

tico/passes/decompose_fake_quantize_tensor_qparams.py

@@ -30,6 +30,7 @@ from torch.ao.quantization.fx._decomposed import quantized_decomposed_lib
 from torch.export import ExportedProgram
 
 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_const_diff_on_pass,
@@ -200,6 +201,7 @@ class DecomposeFakeQuantizeTensorQParams(PassBase):
         modified = False
 
         gm = exported_program.graph_module
+        g = gm.graph
         qd = torch.ops.quantized_decomposed  # type: ignore[return]
         for node in gm.graph.nodes:
             if node.op != "call_function":
@@ -226,17 +228,19 @@ class DecomposeFakeQuantizeTensorQParams(PassBase):
                     **{"dtype": get_quant_type(quant_min, quant_max)},
                 }
                 with gm.graph.inserting_before(node):
-                    quant = gm.graph.call_function(
+                    quant = create_node(
+                        g,
                         qd.quantize_per_tensor.default,
                         args=(tensor, s_value, zp_value, quant_min, quant_max),
                         kwargs=quant_kwargs,
+                        origin=node,
                     )
-                    dequant = gm.graph.call_function(
+                    dequant = create_node(
+                        g,
                         qd.dequantize_per_tensor.default,
                         args=(quant, *quant.args[1:]),
                         kwargs=quant.kwargs,
                     )
-                    # Not set meta for propagating replacing get_item's meta.
                     get_item.replace_all_uses_with(dequant, propagate_meta=True)
                     # If `mask` can be graph output, which prevents `eliminate_dead_code()` from eliminating `mask`.
                     # So, let's remove `mask` from the output.args first.
@@ -267,17 +271,19 @@ class DecomposeFakeQuantizeTensorQParams(PassBase):
                     **{"dtype": get_quant_type(quant_min, quant_max)},
                 }
                 with gm.graph.inserting_before(node):
-                    quant = gm.graph.call_function(
+                    quant = create_node(
+                        g,
                         qd.quantize_per_tensor.default,
                         args=(tensor, s_value, zp_value, quant_min, quant_max),
                         kwargs=quant_kwargs,
+                        origin=node,
                     )
-                    dequant = gm.graph.call_function(
+                    dequant = create_node(
+                        g,
                         qd.dequantize_per_tensor.default,
                         args=(quant, *quant.args[1:]),
                         kwargs=quant.kwargs,
                     )
-                    # Not set meta for propagating replacing get_item's meta.
                     node.replace_all_uses_with(dequant, propagate_meta=True)
                 modified = True