Fix get_control_flow_submodules_list call in debug handle generator pass (#6187)

Summary:

We should be using `current_graph_module` not `ep.graph_module` in this loop, otherwise it can go into an infinite loop.

Reviewed By: Olivia-liu, bingcy

Differential Revision: D64309046

Author: tarun292

exir/passes/debug_handle_generator_pass.py

@@ -51,7 +51,7 @@ def get_control_flow_submodules_list(graph_module):
         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)
+        control_flow_submodules = get_control_flow_submodules_list(current_graph_module)
         queue.extend(control_flow_submodules)
 
     queue = [ep.graph_module]
@@ -61,5 +61,5 @@ def get_control_flow_submodules_list(graph_module):
             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)
+        control_flow_submodules = get_control_flow_submodules_list(current_graph_module)
         queue.extend(control_flow_submodules)

exir/tests/test_passes.py

@@ -1018,16 +1018,13 @@ def forward(
             torch.ones(2, 2),
         )
 
-        graph_module = (
-            to_edge(
-                export(
-                    f,
-                    inputs,
-                )
+        ep = to_edge(
+            export(
+                f,
+                inputs,
             )
-            .exported_program()
-            .graph_module
-        )
+        ).exported_program()
+        graph_module = ep.graph_module
 
         def check_debug_handle_metadata(graph_module: torch.fx.GraphModule) -> None:
             queue = [graph_module]
@@ -1045,6 +1042,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)
 
         # Check debug handle still preserved after ScalarToTensorPass
         ScalarToTensorPass()(graph_module)