tico 0.1.0.dev250609__py3-none-any.whl → 0.1.0.dev250611__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.dev250609"
+__version__ = "0.1.0.dev250611"
 
 MINIMUM_SUPPORTED_VERSION = "2.5.0"
 SECURE_TORCH_VERSION = "2.6.0"

tico/passes/cast_aten_where_arg_type.py

@@ -26,7 +26,8 @@ from tico.utils.trace_decorators import (
     trace_const_diff_on_pass,
     trace_graph_diff_on_pass,
 )
-from tico.utils.utils import set_new_meta_val
+from tico.utils.utils import is_target_node, set_new_meta_val
+from tico.utils.validate_args_kwargs import WhereSelfArgs
 
 
 dtype_ranking = {
@@ -114,69 +115,71 @@ class CastATenWhereArgType(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if node.op == "call_function" and node.target == torch.ops.aten.where.self:
-
-                assert len(node.args) == 3
-                (
-                    _,
-                    result_true,
-                    result_false,
-                ) = node.args  # first argument is not used
-
-                ep = exported_program
-
-                if not (
-                    result_true.name in ep.graph_signature.inputs_to_buffers
-                    and result_false.name in ep.graph_signature.inputs_to_buffers
-                ):
-                    continue
-
-                # Check if they have different data types
-                true_dtype = extract_torch_dtype(result_true)
-                false_dtype = extract_torch_dtype(result_false)
-                if true_dtype == false_dtype:
-                    continue
-
-                node_to_dtype = {result_true: true_dtype, result_false: false_dtype}
-
-                not_to_cast, to_cast = sort_by_dtype(result_true, result_false)
-
-                buf_name_to_data = {name: buf for name, buf in ep.named_buffers()}
-                buf_name = ep.graph_signature.inputs_to_buffers[to_cast.name]
-                buf_data = buf_name_to_data[buf_name]
-
-                assert isinstance(buf_data, torch.Tensor)
-
-                dtype_to_cast = node_to_dtype[not_to_cast]
-
-                if dtype_to_cast == torch.float32:
-                    if not check_if_covered_by_float(buf_data):
-                        raise RuntimeError(
-                            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(
-                        torch.ops.aten._to_copy.default,
-                        args=(to_cast,),
-                        kwargs={"dtype": dtype_to_cast},
+            if not is_target_node(node, torch.ops.aten.where.self):
+                continue
+
+            where_args = WhereSelfArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+            result_true, result_false = where_args.input, where_args.other
+            if not isinstance(result_true, torch.fx.Node) or not isinstance(
+                result_false, torch.fx.Node
+            ):
+                continue
+
+            ep = exported_program
+            assert isinstance(result_true, torch.fx.Node)
+            assert isinstance(result_false, torch.fx.Node)
+            if not (
+                result_true.name in ep.graph_signature.inputs_to_buffers
+                and result_false.name in ep.graph_signature.inputs_to_buffers
+            ):
+                continue
+
+            # Check if they have different data types
+            true_dtype = extract_torch_dtype(result_true)
+            false_dtype = extract_torch_dtype(result_false)
+            if true_dtype == false_dtype:
+                continue
+
+            node_to_dtype = {result_true: true_dtype, result_false: false_dtype}
+
+            not_to_cast, to_cast = sort_by_dtype(result_true, result_false)
+
+            buf_name_to_data = {name: buf for name, buf in ep.named_buffers()}
+            buf_name = ep.graph_signature.inputs_to_buffers[to_cast.name]
+            buf_data = buf_name_to_data[buf_name]
+
+            assert isinstance(buf_data, torch.Tensor)
+
+            dtype_to_cast = node_to_dtype[not_to_cast]
+
+            if dtype_to_cast == torch.float32:
+                if not check_if_covered_by_float(buf_data):
+                    raise RuntimeError(
+                        f"{to_cast.name}({buf_data.dtype}) data range is out of {dtype_to_cast} range"
                     )
-                # set new meta["val"] in advance because we will use it below for checking if type promotion is valid.
-                set_new_meta_val(cast)
-                node.update_arg(node.args.index(to_cast), cast)
-
-                # check if type promotion is valid.
-                node_dtype_ori = extract_torch_dtype(node)
-                set_new_meta_val(node)
-                node_dtype = extract_torch_dtype(node)
-                assert (
-                    node_dtype == node_dtype_ori
-                ), f"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}"
+            with graph_module.graph.inserting_after(to_cast):
+                cast = graph_module.graph.call_function(
+                    torch.ops.aten._to_copy.default,
+                    args=(to_cast,),
+                    kwargs={"dtype": dtype_to_cast},
                 )
-
-                modified = True
+            # set new meta["val"] in advance because we will use it below for checking if type promotion is valid.
+            set_new_meta_val(cast)
+            node.update_arg(node.args.index(to_cast), cast)
+
+            # check if type promotion is valid.
+            node_dtype_ori = extract_torch_dtype(node)
+            set_new_meta_val(node)
+            node_dtype = extract_torch_dtype(node)
+            assert (
+                node_dtype == node_dtype_ori
+            ), f"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}"
+            )
+
+            modified = True
 
         graph.eliminate_dead_code()
         graph.lint()

tico/passes/cast_mixed_type_args.py

@@ -28,7 +28,7 @@ from tico.serialize.circle_mapping import extract_torch_dtype
 from tico.utils import logging
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
-from tico.utils.utils import set_new_meta_val
+from tico.utils.utils import is_target_node, set_new_meta_val
 
 
 ops_to_promote = {
@@ -96,10 +96,7 @@ class CastMixedTypeArgs(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
-                continue
-
-            if node.target not in ops_to_promote:
+            if not is_target_node(node, list(ops_to_promote.keys())):
                 continue
 
             assert len(node.args) == 2

tico/passes/convert_conv1d_to_conv2d.py

@@ -24,6 +24,7 @@ from tico.utils import logging
 from tico.utils.errors import NotYetSupportedError
 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 Conv1DArgs
 
 
@@ -131,17 +132,15 @@ class ConvertConv1dToConv2d(PassBase):
         return modified
 
     def call(self, exported_program: ExportedProgram) -> PassResult:
-        target_conv_op = (torch.ops.aten.conv1d.default, torch.ops.aten.conv1d.padding)
+        target_conv_op = [torch.ops.aten.conv1d.default, torch.ops.aten.conv1d.padding]
 
         graph_module = exported_program.graph_module
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, target_conv_op):
                 continue
 
-            if node.target not in target_conv_op:
-                continue
             modified |= self.convert(exported_program, node)
 
         graph.eliminate_dead_code()

tico/passes/convert_repeat_to_expand_copy.py

@@ -22,6 +22,8 @@ from torch.export import ExportedProgram
 from tico.utils import logging
 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 RepeatArgs
 
 
 @trace_graph_diff_on_pass
@@ -42,17 +44,11 @@ class ConvertRepeatToExpandCopy(PassBase):
         graph = graph_module.graph
         modified = False
         for node in graph.nodes:
-            if not node.op == "call_function":
+            if not is_target_node(node, torch.ops.aten.repeat.default):
                 continue
 
-            if node.target != torch.ops.aten.repeat.default:
-                continue
-
-            assert len(node.args) == 2
-
-            tensor, repeats = node.args
-            assert isinstance(tensor, torch.fx.Node)
-            assert isinstance(repeats, list)
+            reshape_args = RepeatArgs(*node.args, **node.kwargs)  # type: ignore[arg-type]
+            tensor, repeats = reshape_args.input, reshape_args.repeats
 
             tensor_shape: List[int] = [int(dim) for dim in tensor.meta["val"].shape]
 

tico/passes/decompose_addmm.py

@@ -24,7 +24,7 @@ from tico.utils import logging
 from tico.utils.graph import add_placeholder
 from tico.utils.passes import PassBase, PassResult
 from tico.utils.trace_decorators import trace_graph_diff_on_pass
-from tico.utils.utils import set_new_meta_val
+from tico.utils.utils import is_target_node, set_new_meta_val
 from tico.utils.validate_args_kwargs import AddmmArgs
 
 
@@ -66,59 +66,52 @@ class DecomposeAddmm(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if node.op != "call_function":
+            if not is_target_node(node, torch.ops.aten.addmm.default):
                 continue
 
-            if node.target in [
-                torch.ops.aten.addmm.default,
-            ]:
-                args = AddmmArgs(*node.args, **node.kwargs)
-                input = args.input
-                mat1 = args.mat1
-                mat2 = args.mat2
-                beta = args.beta
-                alpha = args.alpha
-
-                with graph.inserting_before(node):
-                    # out = beta * input + alpha * (mat1 @ mat2)
-                    matmul = graph.call_function(
-                        torch.ops.aten.mm.default, (mat1, mat2)
+            args = AddmmArgs(*node.args, **node.kwargs)
+            input = args.input
+            mat1 = args.mat1
+            mat2 = args.mat2
+            beta = args.beta
+            alpha = args.alpha
+
+            with graph.inserting_before(node):
+                # out = beta * input + alpha * (mat1 @ mat2)
+                matmul = graph.call_function(torch.ops.aten.mm.default, (mat1, mat2))
+                set_new_meta_val(matmul)
+
+                if beta == 1:
+                    bias: torch.fx.Node | torch.Tensor = input
+                elif beta == 0:
+                    bias = add_placeholder(
+                        exported_program,
+                        torch.zeros(extract_shape(input)),
+                        f"{node.name}_beta_zeros",
                     )
-                    set_new_meta_val(matmul)
-
-                    if beta == 1:
-                        bias: torch.fx.Node | torch.Tensor = input
-                    elif beta == 0:
-                        bias = add_placeholder(
-                            exported_program,
-                            torch.zeros(extract_shape(input)),
-                            f"{node.name}_beta_zeros",
-                        )
-                    else:
-                        bias = graph.call_function(
-                            torch.ops.aten.mul.Tensor, (input, beta)
-                        )
-
-                    if alpha == 1:
-                        scaled_matmul: torch.fx.Node | torch.Tensor = matmul
-                    elif alpha == 0:
-                        scaled_matmul = add_placeholder(
-                            exported_program,
-                            torch.zeros(extract_shape(matmul)),
-                            f"{node.name}_alpha_zeros",
-                        )
-                    else:
-                        scaled_matmul = graph.call_function(
-                            torch.ops.aten.mul.Tensor, (matmul, alpha)
-                        )
-
-                    result = graph.call_function(
-                        torch.ops.aten.add.Tensor, (bias, scaled_matmul)
+                else:
+                    bias = graph.call_function(torch.ops.aten.mul.Tensor, (input, beta))
+
+                if alpha == 1:
+                    scaled_matmul: torch.fx.Node | torch.Tensor = matmul
+                elif alpha == 0:
+                    scaled_matmul = add_placeholder(
+                        exported_program,
+                        torch.zeros(extract_shape(matmul)),
+                        f"{node.name}_alpha_zeros",
                     )
+                else:
+                    scaled_matmul = graph.call_function(
+                        torch.ops.aten.mul.Tensor, (matmul, alpha)
+                    )
+
+                result = graph.call_function(
+                    torch.ops.aten.add.Tensor, (bias, scaled_matmul)
+                )
 
-                node.replace_all_uses_with(result, propagate_meta=True)
+            node.replace_all_uses_with(result, propagate_meta=True)
 
-                modified = True
+            modified = True
 
         gm.graph.eliminate_dead_code()
         gm.graph.lint()

tico/passes/decompose_batch_norm.py

@@ -32,7 +32,7 @@ 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
+from tico.utils.utils import fill_meta_val, is_target_node
 from tico.utils.validate_args_kwargs import NativeBatchNormLegitNoTrainingArgs
 
 
@@ -87,109 +87,107 @@ class DecomposeBatchNorm(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if node.op != "call_function":
+            if not is_target_node(
+                node, torch.ops.aten._native_batch_norm_legit_no_training.default
+            ):
                 continue
 
-            if node.target in [
-                torch.ops.aten._native_batch_norm_legit_no_training.default,
-            ]:
-                args = NativeBatchNormLegitNoTrainingArgs(*node.args)
-                input_ = args.input
-                weight = args.weight
-                bias = args.bias
-                running_mean = args.running_mean
-                running_var = args.running_var
-                eps = args.eps
-
-                if not running_mean:
-                    raise NotYetSupportedError(
-                        f"running_mean=None is not supported yet"
-                    )
-                if not running_var:
-                    raise NotYetSupportedError(f"running_var=None is not supported yet")
-
-                """
-                Only support the cases generated from torch.nn.BatchNorm2d module,
-                 for which, let's checks if weight and bias are parameters and 
-                 running_mean and running_var are buffers.
-                """
-                if weight and not is_torch_param(weight, exported_program):
-                    continue
-                if bias and not is_torch_param(bias, exported_program):
-                    continue
-                if not is_torch_buffer(running_mean, exported_program):
-                    continue
-                if not is_torch_buffer(running_var, exported_program):
-                    continue
-
-                input_shape = extract_shape(input_)
-                assert len(input_shape) == 4
-                C = input_shape[1]
-
-                weight_value = (
-                    get_torch_param_value(weight, exported_program)
-                    if weight
-                    else torch.tensor([1] * C)
+            args = NativeBatchNormLegitNoTrainingArgs(*node.args)
+            input_ = args.input
+            weight = args.weight
+            bias = args.bias
+            running_mean = args.running_mean
+            running_var = args.running_var
+            eps = args.eps
+
+            if not running_mean:
+                raise NotYetSupportedError(f"running_mean=None is not supported yet")
+            if not running_var:
+                raise NotYetSupportedError(f"running_var=None is not supported yet")
+
+            """
+            Only support the cases generated from torch.nn.BatchNorm2d module,
+             for which, let's checks if weight and bias are parameters and 
+             running_mean and running_var are buffers.
+            """
+            if weight and not is_torch_param(weight, exported_program):
+                continue
+            if bias and not is_torch_param(bias, exported_program):
+                continue
+            if not is_torch_buffer(running_mean, exported_program):
+                continue
+            if not is_torch_buffer(running_var, exported_program):
+                continue
+
+            input_shape = extract_shape(input_)
+            assert len(input_shape) == 4
+            C = input_shape[1]
+
+            weight_value = (
+                get_torch_param_value(weight, exported_program)
+                if weight
+                else torch.tensor([1] * C)
+            )
+            bias_value = (
+                get_torch_param_value(bias, exported_program)
+                if bias
+                else torch.tensor([0] * C)
+            )
+            mean_value = get_torch_buffer_value(running_mean, exported_program)
+            var_value = get_torch_buffer_value(running_var, exported_program)
+
+            assert isinstance(weight_value, torch.Tensor)
+            assert isinstance(bias_value, torch.Tensor)
+            assert isinstance(mean_value, torch.Tensor)
+            assert isinstance(var_value, torch.Tensor)
+
+            assert (
+                weight_value.shape
+                == bias_value.shape
+                == mean_value.shape
+                == var_value.shape
+            )
+            # Calculate constants for mul and add
+            mul_const = weight_value / torch.sqrt(var_value + eps)
+            add_const = bias_value - (mul_const * mean_value)
+            # N, C, H, W
+            assert len(mul_const) == len(add_const) == C
+            # reshape along with channel dimension
+            mul_const = mul_const.view(1, mul_const.shape[0], 1, 1)
+            add_const = add_const.view(1, add_const.shape[0], 1, 1)
+
+            # Placeholder nodes must be the first N nodes in the nodes list of a graph.
+            # Therefore, insert the newly created placeholders at the start of the node list.
+            with exported_program.graph.inserting_before(
+                get_first_user_input(exported_program)
+            ):
+                mul_const_node = add_placeholder(
+                    exported_program,
+                    mul_const,
+                    prefix=f"{node.name}_mul_const",
+                )
+                add_const_node = add_placeholder(
+                    exported_program,
+                    add_const,
+                    prefix=f"{node.name}_add_const",
                 )
-                bias_value = (
-                    get_torch_param_value(bias, exported_program)
-                    if bias
-                    else torch.tensor([0] * C)
+
+            with gm.graph.inserting_before(node):
+                mul = graph.call_function(
+                    torch.ops.aten.mul.Tensor,
+                    args=(input_, mul_const_node),
                 )
-                mean_value = get_torch_buffer_value(running_mean, exported_program)
-                var_value = get_torch_buffer_value(running_var, exported_program)
-
-                assert isinstance(weight_value, torch.Tensor)
-                assert isinstance(bias_value, torch.Tensor)
-                assert isinstance(mean_value, torch.Tensor)
-                assert isinstance(var_value, torch.Tensor)
-
-                assert (
-                    weight_value.shape
-                    == bias_value.shape
-                    == mean_value.shape
-                    == var_value.shape
+                add = graph.call_function(
+                    torch.ops.aten.add.Tensor,
+                    args=(mul, add_const_node),
                 )
-                # Calculate constants for mul and add
-                mul_const = weight_value / torch.sqrt(var_value + eps)
-                add_const = bias_value - (mul_const * mean_value)
-                # N, C, H, W
-                assert len(mul_const) == len(add_const) == C
-                # reshape along with channel dimension
-                mul_const = mul_const.view(1, mul_const.shape[0], 1, 1)
-                add_const = add_const.view(1, add_const.shape[0], 1, 1)
-
-                # Placeholder nodes must be the first N nodes in the nodes list of a graph.
-                # Therefore, insert the newly created placeholders at the start of the node list.
-                with exported_program.graph.inserting_before(
-                    get_first_user_input(exported_program)
-                ):
-                    mul_const_node = add_placeholder(
-                        exported_program,
-                        mul_const,
-                        prefix=f"{node.name}_mul_const",
-                    )
-                    add_const_node = add_placeholder(
-                        exported_program,
-                        add_const,
-                        prefix=f"{node.name}_add_const",
-                    )
-
-                with gm.graph.inserting_before(node):
-                    mul = graph.call_function(
-                        torch.ops.aten.mul.Tensor,
-                        args=(input_, mul_const_node),
-                    )
-                    add = graph.call_function(
-                        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)
-                modified = True
+                # 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
 
         gm.graph.eliminate_dead_code()
         gm.graph.lint()

tico/passes/decompose_fake_quantize.py

@@ -71,9 +71,9 @@ class DecomposeFakeQuantize(PassBase):
         gm = exported_program.graph_module
         qd = torch.ops.quantized_decomposed  # type: ignore[return]
         for node in gm.graph.nodes:
-            if node.op == "call_function" and node.target in [
-                torch.ops.aten.fake_quantize_per_tensor_affine.default
-            ]:
+            if node.op != "call_function":
+                continue
+            if node.target in [torch.ops.aten.fake_quantize_per_tensor_affine.default]:
                 # tensor, scale, zero_p, quant_min, quant_max
                 assert len(node.args) == 5
                 _, _, _, quant_min, quant_max = node.args
@@ -97,9 +97,7 @@ class DecomposeFakeQuantize(PassBase):
                     node.replace_all_uses_with(dequnt, propagate_meta=True)
                 modified = True
 
-            if node.op == "call_function" and node.target in [
-                torch.ops.aten.fake_quantize_per_channel_affine.default
-            ]:
+            if node.target in [torch.ops.aten.fake_quantize_per_channel_affine.default]:
                 fq_args = FakeQuantizePerChannelArgs(*node.args, **node.kwargs)
                 quant_min = fq_args.quant_min
                 quant_max = fq_args.quant_max

tico/passes/decompose_fake_quantize_tensor_qparams.py

@@ -202,9 +202,10 @@ class DecomposeFakeQuantizeTensorQParams(PassBase):
         gm = exported_program.graph_module
         qd = torch.ops.quantized_decomposed  # type: ignore[return]
         for node in gm.graph.nodes:
+            if node.op != "call_function":
+                continue
             if (
-                node.op == "call_function"
-                and node.target
+                node.target
                 == torch.ops.aten._fake_quantize_per_tensor_affine_cachemask_tensor_qparams.default
             ):
                 # tensor, scale, zero_p, fake_quant_enabled, quant_min, quant_max
@@ -247,8 +248,7 @@ class DecomposeFakeQuantizeTensorQParams(PassBase):
                         mask_user.args = ((mask_user.args[0][0],),)
                 modified = True
             if (
-                node.op == "call_function"
-                and node.target
+                node.target
                 == torch.ops.aten.fake_quantize_per_tensor_affine.tensor_qparams
             ):
                 fq_args = FakeQuantizePerTensorTQParamArgs(*node.args, **node.kwargs)

tico/passes/decompose_group_norm.py

@@ -25,6 +25,7 @@ from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 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 NativeGroupNormArgs, NativeLayerNormArgs
 
 
@@ -115,13 +116,13 @@ class DecomposeGroupNorm(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if node.op != "call_function":
-                continue
-
-            if node.target not in [
-                torch.ops.aten.native_layer_norm.default,
-                torch.ops.aten.native_group_norm.default,
-            ]:
+            if not is_target_node(
+                node,
+                [
+                    torch.ops.aten.native_layer_norm.default,
+                    torch.ops.aten.native_group_norm.default,
+                ],
+            ):
                 continue
 
             if node.target == torch.ops.aten.native_layer_norm.default:

tico/passes/decompose_grouped_conv2d.py

@@ -26,6 +26,7 @@ from tico.utils.errors import InvalidArgumentError, NotYetSupportedError
 from tico.utils.graph import add_placeholder
 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 Conv2DArgs
 
 
@@ -88,9 +89,7 @@ class DecomposeGroupedConv2d(PassBase):
         modified = False
 
         for node in graph.nodes:
-            if node.op != "call_function":
-                continue
-            if not node.target in ops.aten.conv2d:
+            if not is_target_node(node, ops.aten.conv2d):
                 continue
 
             args = Conv2DArgs(*node.args)