Update on "propagate debug handle from edge dialect graph back to exported graph"

Using exported graph from torch.export as source of truth for aot intermediate output is our target. Once thing blocking us is the exported graph does not have debug handle, which will not show up in the export flow until DebugHandleGenerationPass, the last step of to_edge(). We need to equip the graph with same debug handle used in ExecuTorch flow.

This diff creates a propagate_back_debug_handle function, which propagate debug handle from edge dialect program back to the exported program while maintain the correctness of operator tracing.

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

[ghstack-poisoned]

Author: Gasoonjia

devtools/etrecord/_etrecord.py

@@ -244,7 +244,7 @@ def generate_etrecord(
     export_graph_id = 0
 
     if exported_program is not None:
-        # If multiple exported programs are provided, only saved forward method
+        # If multiple exported programs are provided, only save forward method
         if isinstance(exported_program, dict) and "forward" in exported_program:
             exported_program = exported_program["forward"]
 

devtools/inspector/_inspector_utils.py

@@ -958,6 +958,8 @@ def _equip_debug_handle(node: torch.fx.Node) -> None:
             node.meta[DEBUG_HANDLE_KEY] = debug_handle_for_removed_node
 
     bfs_trace_with_node_process(exported_program.graph_module, _find_n_match_node)
+
+    # 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):
         return False
 

devtools/inspector/tests/inspector_utils_test.py

@@ -46,7 +46,7 @@
     TimeScale,
 )
 from executorch.devtools.inspector.numerical_comparator import L1Comparator
-from executorch.exir import EdgeCompileConfig, to_edge
+from executorch.exir import to_edge
 from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY
 from torch.export import export