tico 0.1.0.dev250616__py3-none-any.whl → 0.1.0.dev250617__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.dev250617"
 
 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/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
 

tico/passes/decompose_group_norm.py

@@ -23,6 +23,7 @@ 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.utils import is_target_node
@@ -89,24 +90,40 @@ class DecomposeGroupNorm(PassBase):
     def __init__(self):
         super().__init__()
 
-    def _insert_norm(self, graph, tensor, eps):
+    def _insert_norm(self, graph, tensor, eps, origin):
         """
         Insert (tensor - mean) / sqrt(var + eps)) into the graph
           and return the normalized tensor node.
         """
-        mean = graph.call_function(
-            torch.ops.aten.mean.dim, (tensor, [-1]), {"keepdim": True}
+        mean = create_node(
+            graph,
+            torch.ops.aten.mean.dim,
+            (tensor, [-1]),
+            {"keepdim": True},
+            origin=origin,
         )
-        deviation = graph.call_function(torch.ops.aten.sub.Tensor, (tensor, mean))
-        squared = graph.call_function(torch.ops.aten.pow.Tensor_Scalar, (deviation, 2))
-        var = graph.call_function(
-            torch.ops.aten.mean.dim, (squared, [-1]), {"keepdim": True}
+        deviation = create_node(
+            graph, torch.ops.aten.sub.Tensor, (tensor, mean), origin=origin
         )
-        inverse_std = graph.call_function(
+        squared = create_node(
+            graph, torch.ops.aten.pow.Tensor_Scalar, (deviation, 2), origin=origin
+        )
+        var = create_node(
+            graph,
+            torch.ops.aten.mean.dim,
+            (squared, [-1]),
+            {"keepdim": True},
+            origin=origin,
+        )
+        inverse_std = create_node(
+            graph,
             torch.ops.aten.rsqrt.default,
-            (graph.call_function(torch.ops.aten.add.Tensor, (var, eps)),),
+            (create_node(graph, torch.ops.aten.add.Tensor, (var, eps), origin=origin),),
+            origin=origin,
+        )
+        return create_node(
+            graph, torch.ops.aten.mul.Tensor, (deviation, inverse_std), origin=origin
         )
-        return graph.call_function(torch.ops.aten.mul.Tensor, (deviation, inverse_std))
 
     def call(self, exported_program: ExportedProgram) -> PassResult:
         logger = logging.getLogger(__name__)
@@ -178,17 +195,23 @@ class DecomposeGroupNorm(PassBase):
                 # Branch only on whether a reshape is needed; the normalization is shared.
                 if norm_size != x_shape[-1]:
                     # Pack groups so that the last dimension equals norm_size.
-                    packed = graph.call_function(
-                        torch.ops.aten.reshape.default, (x, pack_shape)
+                    packed = create_node(
+                        graph,
+                        torch.ops.aten.reshape.default,
+                        (x, pack_shape),
+                        origin=node,
                     )
-                    normed = self._insert_norm(graph, packed, eps)
+                    normed = self._insert_norm(graph, packed, eps, origin=node)
                     # Restore the original shape after normalization.
-                    layer_norm = graph.call_function(
-                        torch.ops.aten.reshape.default, (normed, x_shape)
+                    layer_norm = create_node(
+                        graph,
+                        torch.ops.aten.reshape.default,
+                        (normed, x_shape),
+                        origin=node,
                     )
                 else:
                     # The input already has norm_size in the last dimension.
-                    layer_norm = self._insert_norm(graph, x, eps)
+                    layer_norm = self._insert_norm(graph, x, eps, origin=node)
 
                 # weight
                 if weight:
@@ -197,13 +220,17 @@ class DecomposeGroupNorm(PassBase):
                         assert weight_shape[0] == C
                         reshape_size = [1] * len(x_shape)
                         reshape_size[1] = C
-                        weight = graph.call_function(
+                        weight = create_node(
+                            graph,
                             torch.ops.aten.view.default,
                             (weight, reshape_size),
+                            origin=node,
                         )
-                    layer_norm = graph.call_function(
+                    layer_norm = create_node(
+                        graph,
                         torch.ops.aten.mul.Tensor,
                         (layer_norm, weight),
+                        origin=node,
                     )
 
                 # bias
@@ -213,15 +240,17 @@ class DecomposeGroupNorm(PassBase):
                         assert bias_shape[0] == C
                         reshape_size = [1] * len(x_shape)
                         reshape_size[1] = C
-                        bias = graph.call_function(
+                        bias = create_node(
+                            graph,
                             torch.ops.aten.view.default,
                             (bias, reshape_size),
+                            origin=node,
                         )
-                    layer_norm = graph.call_function(
+                    layer_norm = create_node(
+                        graph,
                         torch.ops.aten.add.Tensor,
                         (layer_norm, bias),
                     )
-
                 # Reset last node's meta for propagating replacing node's meta.
                 layer_norm.meta = {}
 

tico/passes/decompose_grouped_conv2d.py

@@ -23,7 +23,7 @@ from tico.passes import ops
 from tico.serialize.circle_mapping import extract_shape
 from tico.utils import logging
 from tico.utils.errors import InvalidArgumentError, NotYetSupportedError
-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
@@ -159,19 +159,26 @@ class DecomposeGroupedConv2d(PassBase):
 
                 conv2d_tensors = []
                 for i in range(groups):
-                    sliced_input = graph.call_function(
+                    sliced_input = create_node(
+                        graph,
                         torch.ops.aten.slice.Tensor,
                         (input_, 1, group_size * i, group_size * (i + 1), 1),
+                        origin=node,
                     )
-                    sliced_weight = graph.call_function(
+                    sliced_weight = create_node(
+                        graph,
                         torch.ops.aten.slice.Tensor,
                         (weight, 0, out_group_size * i, out_group_size * (i + 1), 1),
+                        origin=node,
                     )
-                    sliced_bias = graph.call_function(
+                    sliced_bias = create_node(
+                        graph,
                         torch.ops.aten.slice.Tensor,
                         (bias, 0, out_group_size * i, out_group_size * (i + 1), 1),
+                        origin=node,
                     )
-                    conv2d_tensor = graph.call_function(
+                    conv2d_tensor = create_node(
+                        graph,
                         conv2d_op,
                         (
                             sliced_input,
@@ -182,11 +189,12 @@ class DecomposeGroupedConv2d(PassBase):
                             dilation,
                             1,
                         ),
+                        origin=node,
                     )
                     conv2d_tensors.append(conv2d_tensor)
 
-                concat_output = graph.call_function(
-                    torch.ops.aten.cat.default, (conv2d_tensors, 1)
+                concat_output = create_node(
+                    graph, torch.ops.aten.cat.default, (conv2d_tensors, 1)
                 )
 
                 node.replace_all_uses_with(concat_output, propagate_meta=True)