generate debug handle before opeartor decomposition (#11949)

Summary:

This diff update the debug handle generation, from each node in the edge program having a individual debug handle, to all nodes having a same ancestor in export graph sharing a same debug handle, which update the start point of tracing our node transformation from edge graph to exported graph.

Differential Revision: D76860368

Author: Gasoonjia

devtools/inspector/_inspector.py

@@ -654,7 +654,7 @@ def _populate_debugging_related_fields(
 
     def _associate_with_op_graph_nodes(
         self,
-        debug_handle_to_op_node_map: Dict[int, OperatorNode],
+        debug_handle_to_op_node_map: Dict[int, List[OperatorNode]],
     ) -> None:
         """
         Helper function to populate the stack_traces, module_hierarchy and op_types attributes
@@ -672,14 +672,21 @@ def _associate_with_op_graph_nodes(
             debug_handles = [debug_handles]
 
         for handle in debug_handles:
-            node = debug_handle_to_op_node_map.get(handle)
-            # Attach node metadata including stack traces, module hierarchy and op_types to this event
-            if node is not None and (metadata := node.metadata) is not None:
-                self.stack_traces[node.name] = metadata.get("stack_trace")
-                self.module_hierarchy[node.name] = metadata.get("nn_module_stack")
-                if node.op:
-                    # TODO: consider having this as a dict from node.name -> node.op
-                    self.op_types += [node.op]
+            nodes = debug_handle_to_op_node_map.get(handle, None)
+            if nodes is None:
+                continue
+
+            for node in nodes:
+                # Attach node metadata including stack traces, module hierarchy and op_types to this event
+                if node is not None and (metadata := node.metadata) is not None:
+                    if node.name not in self.stack_traces:
+                        self.stack_traces[node.name] = metadata.get("stack_trace")
+                        self.module_hierarchy[node.name] = metadata.get(
+                            "nn_module_stack"
+                        )
+                    if node.op:
+                        # TODO: consider having this as a dict from node.name -> node.op
+                        self.op_types += [node.op]
 
 
 @dataclass

devtools/inspector/_inspector_utils.py

@@ -279,14 +279,18 @@ def gen_graphs_from_etrecord(
     return op_graph_map
 
 
+# One debug handle should only be associated with one node. We are in the middle of migrating debug handle generation
+# from graph after to_edge to graph after torch.export, one every debug handle in exported graph may be associated with multiple nodes in to_edge
+# graph. After fully migration, we should bring the bring type as well as the #node check back.
+# TODO(gasoonjia): recover the return type to Dict[int, List[OperatorNode], reenable the #node check.
 def create_debug_handle_to_op_node_mapping(
     op_graph: OperatorGraph,
-) -> Dict[int, OperatorNode]:
+) -> Dict[int, List[OperatorNode]]:
     """
     Recursive function to traverse all the operator graph nodes of input op_graph and build a mapping
     from each debug handle to the operator node that contains the debug handle in its metadata.
     """
-    debug_handle_to_op_node_map: Dict[int, OperatorNode] = {}
+    debug_handle_to_op_node_map: Dict[int, List[OperatorNode]] = {}
 
     # Recursively searches through the metadata of nodes
     def _extract_debug_handles(graph: OperatorGraph):
@@ -296,14 +300,13 @@ def _extract_debug_handles(graph: OperatorGraph):
             if isinstance(element, OperatorNode) and element.metadata is not None:
                 metadata = element.metadata
                 debug_handle = metadata.get("debug_handle")
-                if debug_handle is not None:
-                    existing_entry = debug_handle_to_op_node_map.get(debug_handle)
-                    if existing_entry is not None:
-                        raise ValueError(
-                            f"Duplicated debug handle {str(debug_handle)} shared between {element.name} and {existing_entry.name}. "
-                            "No two op nodes of the same graph should have the same debug handle."
-                        )
-                    debug_handle_to_op_node_map[debug_handle] = element
+                if debug_handle is None:
+                    continue
+
+                if debug_handle not in debug_handle_to_op_node_map:
+                    debug_handle_to_op_node_map[debug_handle] = []
+
+                debug_handle_to_op_node_map[debug_handle].append(element)
 
     # Start traversing
     _extract_debug_handles(op_graph)

exir/passes/debug_handle_generator_pass.py

@@ -4,10 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from typing import Dict
+
 from executorch.exir.graph_module import bfs_trace_with_node_process
 from executorch.exir.pass_base import ExportPass
 from torch.export import ExportedProgram
-from torch.fx import GraphModule
+from torch.fx import GraphModule, Node
 from torch.fx.passes.infra.pass_base import PassResult
 
 
@@ -17,12 +19,50 @@ def call(self, graph_module: GraphModule) -> PassResult:
         to executorch backend, that has a canonical set of quantized operators
         """
 
-        index = 1
+        FROM_NODE_KEY = "from_node"
+        DEBUG_HANDLE_KEY = "debug_handle"
+
+        source_node_to_debug_handle: Dict[str, int] = {}
+
+        def _get_greatest_ancestor_source_node(node: Node) -> str:
+            """ Get the source of the greatest ancestor node of the given node. The source
+            here means the name of the node concated with the id the graph it belongs to.
+            For example, if the node transformation is node a -> b -> c, then the greatest
+            ancestor node of c is a.
+            """
+
+            node_source = node.meta[FROM_NODE_KEY]
+            node_source = node_source[-1]
+
+            while len(node_source.from_node) > 0:
+                node_source = node_source.from_node[-1]
+
+            return node_source.name + str(node_source.graph_id)
+
+        def _extract_debug_handles_from_node(node: Node) -> None:
+            """
+            Generate a debug handle based on node's oldest ancestor node's name
+            and graph id, or return None if the node does not need to be traced.
+            """
+
+            if node.op == "placeholder" or node.op == "output":
+                # placeholder and output nodes don't have debug handle
+                return
+
+            assert (
+                FROM_NODE_KEY in node.meta
+            ), f"Node {node} does not have meta key {FROM_NODE_KEY}"
+
+            source_node = _get_greatest_ancestor_source_node(node)
+
+            debug_handle = (
+                len(source_node_to_debug_handle) + 1
+                if source_node not in source_node_to_debug_handle
+                else source_node_to_debug_handle[source_node]
+            )
+            source_node_to_debug_handle[source_node] = debug_handle
 
-        def _extract_debug_handles_from_node(node):
-            nonlocal index
-            node.meta["debug_handle"] = index
-            index += 1
+            node.meta[DEBUG_HANDLE_KEY] = debug_handle
 
         bfs_trace_with_node_process(graph_module, _extract_debug_handles_from_node)
 

exir/tests/test_passes.py

@@ -939,7 +939,7 @@ def check_debug_handle_metadata(graph_module: torch.fx.GraphModule) -> None:
 
         DebugHandleGeneratorPass()(graph_module)
         check_debug_handle_metadata(graph_module)
-        generate_missing_debug_handles(ep)
+        # generate_missing_debug_handles(ep)
 
         # Check debug handle still preserved after ScalarToTensorPass
         ScalarToTensorPass()(graph_module)