Add helper method to generate missing debug handles

Summary: Helper function to generate missing debug handles on nodes, which is usually needed when graph transforms are done and new nodes are inserted.

Reviewed By: Vysarat

Differential Revision: D63913905

Author: tarun292

exir/passes/debug_handle_generator_pass.py

@@ -6,6 +6,7 @@
 
 from executorch.exir.graph_module import get_control_flow_submodules
 from executorch.exir.pass_base import ExportPass
+from torch.export import ExportedProgram
 from torch.fx import GraphModule
 from torch.fx.passes.infra.pass_base import PassResult
 
@@ -17,7 +18,7 @@ def call(self, graph_module: GraphModule) -> PassResult:
         """
 
         queue = [graph_module]
-        index = 0
+        index = 1
         # bfs to traverse all modules including control flow submodules to attached debug handle id
         while queue:
             current_graph_module = queue.pop(0)
@@ -30,3 +31,35 @@ def call(self, graph_module: GraphModule) -> PassResult:
             ]
             queue.extend(control_flow_submodules)
         return PassResult(graph_module, True)
+
+
+def generate_missing_debug_handles(ep: ExportedProgram):
+    """
+    This pass is used to generate missing debug handles for the graph module and its submodules.
+    """
+
+    def get_control_flow_submodules_list(graph_module):
+        return [
+            submodule for _, submodule, _ in get_control_flow_submodules(graph_module)
+        ]
+
+    max_handle = 0
+    queue = [ep.graph_module]
+
+    while queue:
+        current_graph_module = queue.pop(0)
+        for node in current_graph_module.graph.nodes:
+            if "debug_handle" in node.meta:
+                max_handle = max(max_handle, node.meta["debug_handle"])
+        control_flow_submodules = get_control_flow_submodules_list(ep.graph_module)
+        queue.extend(control_flow_submodules)
+
+    queue = [ep.graph_module]
+    while queue:
+        current_graph_module = queue.pop(0)
+        for node in current_graph_module.graph.nodes:
+            if node.meta.get("debug_handle", 0) in (0, None):
+                node.meta["debug_handle"] = max_handle + 1
+                max_handle += 1
+        control_flow_submodules = get_control_flow_submodules_list(ep.graph_module)
+        queue.extend(control_flow_submodules)

exir/tests/test_passes.py

@@ -33,7 +33,10 @@
     ToOutVarPass,
 )
 from executorch.exir.passes.constant_prop_pass import constant_prop_pass
-from executorch.exir.passes.debug_handle_generator_pass import DebugHandleGeneratorPass
+from executorch.exir.passes.debug_handle_generator_pass import (
+    DebugHandleGeneratorPass,
+    generate_missing_debug_handles,
+)
 from executorch.exir.passes.insert_write_back_for_buffers_pass import (
     insert_write_back_for_buffers_pass,
 )
@@ -949,13 +952,28 @@ def test_debug_handle_generator_pass(self) -> None:
             .exported_program()
             .graph_module
         )
-        DebugHandleGeneratorPass()(graph_module)
         for node in graph_module.graph.nodes:
             self.assertIn("debug_handle", node.meta)
         ScalarToTensorPass()(graph_module)
         for node in graph_module.graph.nodes:
             self.assertIn("debug_handle", node.meta)
 
+    def test_generate_missing_debug_handles(self) -> None:
+        eager_model = MLP(2, output_size=4)
+        inputs = eager_model.get_random_inputs()
+
+        ep = to_edge(
+            export(
+                eager_model,
+                inputs,
+            )
+        ).exported_program()
+
+        list(ep.graph.nodes)[0].meta.pop("debug_handle")
+        self.assertTrue(list(ep.graph.nodes)[0].meta.get("debug_handle") is None)
+        generate_missing_debug_handles(ep)
+        self.assertTrue(list(ep.graph.nodes)[0].meta.get("debug_handle") is not None)
+
     def test_debug_handle_generator_pass_with_control_flow(self) -> None:
         def true_nested(y: torch.Tensor) -> torch.Tensor:
             y = y + y