support back propagate debug handle to arbitrary ancestor export graph

Currently propagate_back_debug_handle function can only support propagating debug handle back to the greatest ancestor export graph.
This diff update algo to support every possible ancestor export graph on the flow.

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

ghstack-source-id: 296756855
Pull Request resolved: #12580

Author: Gasoonjia

devtools/inspector/_inspector_utils.py

@@ -11,7 +11,7 @@
 from collections.abc import Sequence
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, Dict, IO, List, Mapping, Optional, Tuple, TypeAlias, Union
+from typing import Any, Dict, IO, List, Mapping, Optional, Set, Tuple, TypeAlias, Union
 
 import executorch.devtools.etdump.schema_flatcc as flatcc
 
@@ -37,6 +37,7 @@
 
 from executorch.exir.debug_handle_utils import (
     DEBUG_HANDLE_KEY,
+    FROM_NODE_KEY,
     get_greatest_ancestor_node_identifier,
     UNSET_DEBUG_HANDLE,
 )
@@ -46,6 +47,7 @@
 from tabulate import tabulate
 
 from torch.export import ExportedProgram
+from torch.fx import Node
 
 FORWARD = "forward"
 EDGE_DIALECT_GRAPH_KEY = "edge_dialect_graph_module"
@@ -936,6 +938,44 @@ def compare_intermediate_outputs(a: Any, b: Any, comparator) -> List[float]:
         )
 
 
+def get_ancestor_node_identifiers(node: Node) -> List[str]:
+    """Get the identifier of the ancestor node of the given node, with the graph id the ancestor node lives in.
+
+    The identifier is the concatenation of the node name and graph id of the
+    greatest ancestor node, where the graph id is the unique id for every graph
+    module in the export flow and node name is unique within the same graph module.
+
+    Returns: the identifiers of all its ancestor nodes
+    """
+
+    node_source = node.meta[FROM_NODE_KEY]
+    node_source = node_source[-1]
+    ancestor_node_ids: List[str] = [f"{node_source.name}.{str(node_source.graph_id)}"]
+
+    while len(node_source.from_node) > 0:
+        node_source = node_source.from_node[-1]
+        ancestor_node_ids.append(f"{node_source.name}.{str(node_source.graph_id)}")
+
+    return ancestor_node_ids
+
+
+def get_parent_node_identifier(node: Node) -> Optional[str]:
+    """Get the identifier of the parent node of the given node, with the graph id the parent node lives in.
+
+    The identifier is the concatenation of the node name and graph id of the
+    greatest parent node, where the graph id is the unique id for every graph
+    module in the export flow and node name is unique within the same graph module.
+
+    Returns: the identifier of the parent node, or None if can not find the parent
+    """
+
+    if FROM_NODE_KEY not in node.meta:
+        return None
+
+    node_source = node.meta[FROM_NODE_KEY][-1]
+    return f"{node_source.name}.{str(node_source.graph_id)}"
+
+
 def propagate_back_debug_handle(
     exported_program: ExportedProgram,
     exported_program_graph_id: int,
@@ -953,47 +993,78 @@ def propagate_back_debug_handle(
     Then debug handle of op1 should be same as op1_0, and debug handle of op3 should be same as op3_0 and op3_1.
     The debug handle of op2 will be UNSET_DEBUG_HANDLE for further skipping.
 
-    Return: True if:
-        a. every debug handle in the edge dialect program has a corresponding node in the exported program
-        b. the exported program is the greatest ancestor of the edge dialect program
-
-    Otherwise, return False.
+    Return: True if every debug handle in the edge dialect program has a corresponding node in the exported program, otherwise, return False.
     """
 
-    # 1. set up a mapping from debug handle to identifier of export program's node
+    # 1. set up a mapping from identifier of every possible ancestor node id to debug handle
     # using edge dialect program nodes' debug handles and from_node info
-    export_graph_node_id_to_debug_handle = {
-        get_greatest_ancestor_node_identifier(node): node.meta[DEBUG_HANDLE_KEY]
-        for node in edge_dialect_program.graph.nodes
-        if node.op not in ("placeholder", "output")
-    }
-
-    # 2. equip debug handle to the exported program's nodes using the mapping
-    # number of nodes in the exported program that have matched entry in export_graph_node_id_to_debug_handle
-    n_matched_node = 0
+    ancestors_node_id_to_debug_handle: Dict[str, int] = {}
 
-    def _find_n_match_node(node: torch.fx.Node) -> None:
-        nonlocal n_matched_node
-        if node.name in ("output", "placeholder"):
+    def _extract_node_id_to_debug_handle(node: Node) -> None:
+        nonlocal ancestors_node_id_to_debug_handle
+        if node.op in ("placeholder", "output"):
             return
-        node_id = f"{node.name}.{exported_program_graph_id}"
-        if node_id in export_graph_node_id_to_debug_handle:
-            n_matched_node += 1
+        for ancestor_node_id in get_ancestor_node_identifiers(node):
+            if ancestor_node_id not in ancestors_node_id_to_debug_handle:
+                ancestors_node_id_to_debug_handle[ancestor_node_id] = node.meta[
+                    DEBUG_HANDLE_KEY
+                ]
+            else:
+                assert (
+                    ancestors_node_id_to_debug_handle[ancestor_node_id]
+                    == node.meta[DEBUG_HANDLE_KEY]
+                )
+
+    bfs_trace_with_node_process(
+        edge_dialect_program.graph_module, _extract_node_id_to_debug_handle
+    )
 
-    def _equip_debug_handle(node: torch.fx.Node) -> None:
-        if node.name in ("output", "placeholder"):
+    # 2. verify if every debug handle in the edge dialect program has a corresponding node in the exported program
+    matched_debug_handes: Set[int] = set()
+
+    def _find_n_match_node(node: Node) -> None:
+        nonlocal matched_debug_handes
+        if node.op in ("output", "placeholder"):
             return
         node_id = f"{node.name}.{exported_program_graph_id}"
-        if node_id in export_graph_node_id_to_debug_handle:
-            node.meta[DEBUG_HANDLE_KEY] = export_graph_node_id_to_debug_handle[node_id]
-        else:
-            node.meta[DEBUG_HANDLE_KEY] = UNSET_DEBUG_HANDLE
+        parent_node_id = get_parent_node_identifier(node)
+        if node_id in ancestors_node_id_to_debug_handle:
+            matched_debug_handes.add(ancestors_node_id_to_debug_handle[node_id])
+        elif parent_node_id and parent_node_id in ancestors_node_id_to_debug_handle:
+            matched_debug_handes.add(ancestors_node_id_to_debug_handle[parent_node_id])
 
     bfs_trace_with_node_process(exported_program.graph_module, _find_n_match_node)
 
+    graph_matched = True
+
+    def _check_graph_match(node: Node) -> None:
+        nonlocal graph_matched
+        if node.op in ("output", "placeholder"):
+            return
+
+        if node.meta[DEBUG_HANDLE_KEY] not in matched_debug_handes:
+            graph_matched = False
+
+    bfs_trace_with_node_process(edge_dialect_program.graph_module, _check_graph_match)
+
     # if any node in the edge dialect program has no corresponding node in the exported program, match failed
-    if n_matched_node != len(export_graph_node_id_to_debug_handle):
+    if not graph_matched:
         return False
 
+    # 3. propagate debug handle from edge dialect program back to the exported program while maintain the correctness of operator tracing
+    def _equip_debug_handle(node: Node) -> None:
+        if node.op in ("output", "placeholder"):
+            return
+        node_id = f"{node.name}.{exported_program_graph_id}"
+        parent_node_id = get_parent_node_identifier(node)
+        if node_id in ancestors_node_id_to_debug_handle:
+            node.meta[DEBUG_HANDLE_KEY] = ancestors_node_id_to_debug_handle[node_id]
+        elif parent_node_id and parent_node_id in ancestors_node_id_to_debug_handle:
+            node.meta[DEBUG_HANDLE_KEY] = ancestors_node_id_to_debug_handle[
+                parent_node_id
+            ]
+        else:
+            node.meta[DEBUG_HANDLE_KEY] = UNSET_DEBUG_HANDLE
+
     bfs_trace_with_node_process(exported_program.graph_module, _equip_debug_handle)
     return True

devtools/inspector/tests/inspector_utils_test.py

@@ -654,6 +654,95 @@ def test_equip_debug_handle_to_export_program_success(self):
             exported_program_debug_handles[0], edge_dialect_program_debug_handles[1]
         )
 
+    def test_equip_debug_handle_to_strict_export_program_success(self):
+        """Test that propagate_back_debug_handle returns True and properly equips debug handles."""
+        # Create a test model
+        model = models.FeedForwardBlock(5, 10)
+        inputs = (torch.rand(5, 5),)
+
+        # Export the model
+        exported_program = export(model, inputs, strict=True)
+        export_graph_id = id(exported_program.graph)
+
+        # Convert to edge dialect
+        edge_dialect_program = to_edge(exported_program).exported_program()
+
+        # Call propagate_back_debug_handle
+        result = propagate_back_debug_handle(
+            exported_program, export_graph_id, edge_dialect_program
+        )
+
+        self.assertTrue(result)
+
+        # Check that debug handles are properly equipped in the exported program
+        exported_program_debug_handles = []
+        for node in exported_program.graph.nodes:
+            if node.op not in ("placeholder", "output"):
+                self.assertIn(DEBUG_HANDLE_KEY, node.meta)
+                self.assertIsNotNone(node.meta[DEBUG_HANDLE_KEY])
+                exported_program_debug_handles.append(node.meta[DEBUG_HANDLE_KEY])
+
+        edge_dialect_program_debug_handles = []
+        for node in edge_dialect_program.graph.nodes:
+            if node.op not in ("placeholder", "output"):
+                self.assertIn(DEBUG_HANDLE_KEY, node.meta)
+                self.assertIsNotNone(node.meta[DEBUG_HANDLE_KEY])
+                edge_dialect_program_debug_handles.append(node.meta[DEBUG_HANDLE_KEY])
+
+        # The 0th operator in the exported program (layer_norm) has been decomposed into 0th and 1st ops in edge dialect graph (native_layer_norm and getitem)
+        # So they should have the same debug handle
+        self.assertEqual(
+            exported_program_debug_handles[0], edge_dialect_program_debug_handles[0]
+        )
+        self.assertEqual(
+            exported_program_debug_handles[0], edge_dialect_program_debug_handles[1]
+        )
+
+    def test_equip_debug_handle_to_reexport_program_success(self):
+        """Test that propagate_back_debug_handle returns True and properly equips debug handles."""
+        # Create a test model
+        model = models.FeedForwardBlock(5, 10)
+        inputs = (torch.rand(5, 5),)
+
+        # Export the model
+        init_export_program = export(model, inputs)
+        exported_program = export(init_export_program.module(), inputs)
+        export_graph_id = id(exported_program.graph)
+
+        # Convert to edge dialect
+        edge_dialect_program = to_edge(exported_program).exported_program()
+
+        # Call propagate_back_debug_handle
+        result = propagate_back_debug_handle(
+            exported_program, export_graph_id, edge_dialect_program
+        )
+
+        self.assertTrue(result)
+
+        # Check that debug handles are properly equipped in the exported program
+        exported_program_debug_handles = []
+        for node in exported_program.graph.nodes:
+            if node.op not in ("placeholder", "output"):
+                self.assertIn(DEBUG_HANDLE_KEY, node.meta)
+                self.assertIsNotNone(node.meta[DEBUG_HANDLE_KEY])
+                exported_program_debug_handles.append(node.meta[DEBUG_HANDLE_KEY])
+
+        edge_dialect_program_debug_handles = []
+        for node in edge_dialect_program.graph.nodes:
+            if node.op not in ("placeholder", "output"):
+                self.assertIn(DEBUG_HANDLE_KEY, node.meta)
+                self.assertIsNotNone(node.meta[DEBUG_HANDLE_KEY])
+                edge_dialect_program_debug_handles.append(node.meta[DEBUG_HANDLE_KEY])
+
+        # The 0th operator in the exported program (layer_norm) has been decomposed into 0th and 1st ops in edge dialect graph (native_layer_norm and getitem)
+        # So they should have the same debug handle
+        self.assertEqual(
+            exported_program_debug_handles[0], edge_dialect_program_debug_handles[0]
+        )
+        self.assertEqual(
+            exported_program_debug_handles[0], edge_dialect_program_debug_handles[1]
+        )
+
     def test_equip_debug_handle_to_export_program_failure(self):
         """Test that propagate_back_debug_handle returns False when there's a mismatch."""
         # Create a test model