Support get_attr nodes

Author: kevinwuTT

torch_ttnn/backend.py

@@ -106,7 +106,10 @@ def aten_backend(
 
     # Save aten graph if requested
     if options.gen_op_accuracy_tests:
-        option._aten_fx_graphs.append(copy.deepcopy(gm.graph))
+        # Will this hamper memory usage?
+        graph_copy = copy.deepcopy(gm.graph)
+        graph_copy.owning_module = gm
+        option._aten_fx_graphs.append(graph_copy)
 
     # Save the number of aten ops before compilation
     if option.metrics_path:
@@ -237,7 +240,7 @@ def ttnn_backend(
     options: TorchTtnnOption = None,
 ) -> torch.fx.GraphModule:
     # Save all parameters and inputs if requested
-    if options.gen_op_accuracy_tests:
+    if options.gen_op_accuracy_tests and options._all_inputs is None:
         options._all_inputs = generate_op_accuracy_tests.generate_flat_args(gm, example_inputs)
 
     tracer_option = options.tracer_option

torch_ttnn/generate_op_accuracy_tests.py

@@ -43,7 +43,7 @@ def get_opname(node):
 # rename node names because some wrapper or built-in functions have the same name
 def rename_nodes(graph, prefix):
     for node in graph.nodes:
-        if node.op != "placeholder" and node.op != "output":
+        if node.op != "placeholder" and node.op != "output" and node.op != "get_attr":
             # simplify or put this in a new function
             opname = get_opname(node)
             if not opname.startswith("aten.") and not opname.startswith("ttnn."):
@@ -139,7 +139,14 @@ def rename_input_args_from_graph_break(output_nodes, node):
 
 
 def compute_key(node):
-    return str(node.meta["seq_nr"]) + node.meta["original_aten"]._name + str(node.meta["val"])
+    if "tensor_meta" in node.meta:
+        tensor_meta = node.meta["tensor_meta"]
+    else:
+        tensor_meta = node.meta["val"]
+        # Workaround for layer_norm and other ops that has a list for "val"
+        if isinstance(tensor_meta, tuple):
+            tensor_meta = node.meta["val"][0]
+    return str(node.meta["seq_nr"]) + node.meta["original_aten"]._name + str(tensor_meta)
 
 
 def map_meta_to_aten_node(aten_graph):

torch_ttnn/passes/constant_folding_pass.py

@@ -77,7 +77,10 @@ def _evaluate_node(self, gm: torch.fx.GraphModule, node):
     def _replace_with_constant(self, gm: torch.fx.GraphModule, node, value):
         with gm.graph.inserting_before(node):
             new_node = gm.graph.create_node("get_attr", target=f"_folded_{node.name}", args=(), kwargs=None)
-            new_node.meta = node.meta
+            # Copying all of the meta causes some ops to be missing
+            new_node.meta["seq_nr"] = node.meta["seq_nr"]
+            new_node.meta["original_aten"] = node.meta["original_aten"]
+            new_node.meta["tensor_meta"] = node.meta["tensor_meta"]
 
         gm.register_parameter(f"_folded_{node.name}", torch.nn.Parameter(value, requires_grad=False))