make node name consistent before and after serde

Pull Request resolved: #12531


Node name consistency is important for et because we need to trace the node flow by using node name + graph id + from_node info in the deserialized graph in inspector.

Currently we can not have a consistency experience because we didn't record node name as part of schema, but instead, leverging the output info, which didn't record the node name for multiple output node.

This diff makes node name as part of schema to ensure the consistency. Also create tests for coverage.
ghstack-source-id: 296600899

Differential Revision: [D78380855](https://our.internmc.facebook.com/intern/diff/D78380855/)

Author: Gasoonjia

exir/serde/export_serialize.py

@@ -504,6 +504,7 @@ def handle_call_function(self, node: torch.fx.Node):
             assert len(node.kwargs) == 0
             meta_val = node.meta["val"]
             ex_node = Node(
+                name=node.name,
                 target=self.serialize_operator(node.target),
                 inputs=self.serialize_sym_op_inputs(node.target, node.args),
                 outputs=[
@@ -517,6 +518,7 @@ def handle_call_function(self, node: torch.fx.Node):
             assert len(node.kwargs) == 0
             meta_val = node.meta["val"]
             ex_node = Node(
+                name=node.name,
                 target=self.serialize_operator(node.target),
                 inputs=self.serialize_sym_op_inputs(node.target, node.args),
                 outputs=[
@@ -528,6 +530,7 @@ def handle_call_function(self, node: torch.fx.Node):
             )
         elif isinstance(node.target, torch._ops.OpOverload):
             ex_node = Node(
+                name=node.name,
                 target=self.serialize_operator(node.target),
                 inputs=self.serialize_inputs(node.target, node.args, node.kwargs),
                 outputs=self.serialize_outputs(node),
@@ -536,6 +539,7 @@ def handle_call_function(self, node: torch.fx.Node):
             )
         elif isinstance(node.target, torch._ops.HigherOrderOperator):
             ex_node = Node(
+                name=node.name,
                 target=self.serialize_operator(node.target),
                 inputs=self.serialize_hoo_inputs(node.args, node.kwargs),
                 outputs=self.serialize_hoo_outputs(node),
@@ -1658,7 +1662,7 @@ def deserialize_graph(self, serialized_graph: Graph) -> torch.fx.Graph:
 
     def deserialize_node(self, serialized_node: Node, target: Callable) -> None:
         if target in _SYM_BOOL_OPS or target in _SYM_INT_OPS:
-            name = serialized_node.outputs[0].value.as_name
+            name = serialized_node.name
             args = self.deserialize_sym_op_inputs(serialized_node.inputs)
 
             fx_node = self.graph.create_node("call_function", target, args, {}, name)
@@ -1671,12 +1675,7 @@ def deserialize_node(self, serialized_node: Node, target: Callable) -> None:
             # have names that are consistent with serialized.
             #
             # HOPs don't have schema yet, just check the output lengths and as_tensor attribute
-            name = (
-                serialized_node.outputs[0].as_tensor.name
-                if len(serialized_node.outputs) == 1
-                and hasattr(serialized_node.outputs[0], "as_tensor")
-                else None
-            )
+            name = serialized_node.name
             fx_node = self.graph.create_node(
                 "call_function", target, args, kwargs, name
             )
@@ -1687,11 +1686,9 @@ def deserialize_node(self, serialized_node: Node, target: Callable) -> None:
             # For convenience: if this node returns a single tensor, name the
             # newly-created node after it. This ensures that these tensor values
             # have names that are consistent with serialized.
-            name = (
-                serialized_node.outputs[0].as_tensor.name
-                if _is_single_tensor_return(target)
-                else None  # FX will generate a name for us.
-            )
+
+            name = serialized_node.name
+
             args, kwargs = self.deserialize_inputs(target, serialized_node)
             fx_node = self.graph.create_node(
                 "call_function", target, args, kwargs, name

exir/serde/schema.py

@@ -195,6 +195,7 @@ class NamedArgument:
 
 @dataclass
 class Node:
+    name: str
     target: str
     inputs: List[NamedArgument]
     outputs: List[Argument]

exir/serde/serialize.py

@@ -89,6 +89,7 @@ def handle_call_function(self, node: torch.fx.Node) -> None:
 
         if node.target is memory.alloc:
             ex_node = schema.Node(
+                name=node.name,
                 target="memory.alloc",
                 inputs=self.serialize_alloc_inputs(node.args),
                 outputs=self.serialize_arbitrary_outputs(node),
@@ -99,6 +100,7 @@ def handle_call_function(self, node: torch.fx.Node) -> None:
         elif isinstance(node.target, EdgeOpOverload):
             assert node.target._op is not None
             ex_node = schema.Node(
+                name=node.name,
                 target=self.serialize_operator(node.target),
                 # pyre-ignore Undefined attribute [16]: Item `typing.Callable` of
                 # `typing.Union[typing.Callable[..., typing.Any], str]` has no attribute `_op`.
@@ -111,6 +113,7 @@ def handle_call_function(self, node: torch.fx.Node) -> None:
             return
         elif node.target is delegate.executorch_call_delegate:
             ex_node = schema.Node(
+                name=node.name,
                 target=self.serialize_operator(node.target),
                 inputs=self.serialize_call_delegate_inputs(node.args),
                 outputs=self.serialize_arbitrary_outputs(node),

exir/tests/test_serde.py

@@ -42,6 +42,7 @@ def check_ep(
         ep1: TorchExportedProgram,
         ep2: TorchExportedProgram,
         inputs: Tuple[exir.Value, ...],
+        compare_closeness: bool = False,
     ) -> None:
         """
         Checks if two graphs are equivalent
@@ -55,15 +56,40 @@ def check_ep(
         for orig, loaded in zip(flat_orig_outputs, flat_loaded_outputs, strict=True):
             self.assertTrue(torch.allclose(orig, loaded))
 
+        if compare_closeness:
+            self.assertEqual(len(ep1.graph.nodes), len(ep2.graph.nodes))
+            for node_a, node_b in zip(ep1.graph.nodes, ep2.graph.nodes):
+                self.assertEqual(node_a.target, node_b.target)
+                self.assertEqual(node_a.name, node_b.name)
+                self.assertEqual(node_a.type, node_b.type)
+                self.assertEqual(node_a.op, node_b.op)
+                if node_a.op != "call_function":
+                    continue
+
+                self.assertEqual(
+                    node_a.meta.get("debug_handle"), node_b.meta.get("debug_handle")
+                )
+                from_node_a = node_a.meta.get("from_node")
+                from_node_b = node_b.meta.get("from_node")
+
+                if from_node_a is None:
+                    self.assertIsNone(from_node_b)
+                else:
+                    self.assertIsNotNone(from_node_b)
+                    for node_source_a, node_source_b in zip(from_node_a, from_node_b):
+                        self.assertEqual(
+                            node_source_a.to_dict(), node_source_b.to_dict()
+                        )
+
     # pyre-ignore
     def check_serde(self, m, inputs, check_executorch=True) -> None:
         aten = export(m, inputs, strict=True)
         aten_new = deserialize(serialize(aten))
-        self.check_ep(aten, aten_new, inputs)
+        self.check_ep(aten, aten_new, inputs, compare_closeness=True)
 
         edge = to_edge(aten)
         edge_new = deserialize(serialize(edge.exported_program()))
-        self.check_ep(edge.exported_program(), edge_new, inputs)
+        self.check_ep(edge.exported_program(), edge_new, inputs, compare_closeness=True)
 
         buffer = io.BytesIO()
         exir.save(edge.exported_program(), buffer)