Introduce create_node helper (#147)

This commit introduces create_node helper for keeping some meta information.

TICO-DCO-1.0-Signed-off-by: seongwoo <mhs4670go@naver.com>

Author: mhs4670go

tico/experimental/quantization/passes/insert_quantize_on_dtype_mismatch.py

@@ -24,6 +24,7 @@
 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 @@ def _insert_quantize_op_before(node, inp):
 
             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 @@ def _insert_quantize_op_after(node):
             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 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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
                         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 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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
 
             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 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 @@ def convert(self, exported_program: ExportedProgram, node: torch.fx.Node) -> boo
             )
 
         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 @@ def convert(self, exported_program: ExportedProgram, node: torch.fx.Node) -> boo
             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 @@ def convert(self, exported_program: ExportedProgram, node: torch.fx.Node) -> boo
                     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 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 @@ def convert(node, input):
             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 @@
 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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
             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 @@
 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 @@ def convert(self, exported_program, node):
         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 @@ def convert(self, exported_program, node):
         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 @@ def convert(self, exported_program, node):
         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 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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
 
             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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
                         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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
                         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.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.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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
                 )
 
             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.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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
         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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
                     **{"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 @@ def call(self, exported_program: ExportedProgram) -> PassResult:
                     **{"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