Attach debug handle to control flow modules

Summary: Previously we only add debug handle in the top level graph module but not the ones in control flow. This diff handles the control flow modules too.

Reviewed By: tarun292

Differential Revision: D48651108

fbshipit-source-id: ec08fbbff98786433c454d1f837710f82ab389e5

Author: cccclai

exir/passes/TARGETS

@@ -245,6 +245,7 @@ python_library(
     ],
     deps = [
         "//caffe2:torch",
+        "//executorch/exir:graph_module",
         "//executorch/exir:pass_base",
     ],
 )

exir/passes/debug_handle_generator_pass.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from executorch.exir.graph_module import get_control_flow_submodules
 from executorch.exir.pass_base import ExportPass
 from torch.fx import GraphModule
 from torch.fx.passes.infra.pass_base import PassResult
@@ -14,6 +15,18 @@ def call(self, graph_module: GraphModule) -> PassResult:
         """Lower a quantized reference model (with reference quantized operator patterns)
         to executorch backend, that has a canonical set of quantized operators
         """
-        for index, node in enumerate(graph_module.graph.nodes):
-            node.meta["debug_handle"] = index
+
+        queue = [graph_module]
+        index = 0
+        # bfs to traverse all modules including control flow submodules to attached debug handle id
+        while queue:
+            current_graph_module = queue.pop(0)
+            for node in current_graph_module.graph.nodes:
+                node.meta["debug_handle"] = index
+                index += 1
+            control_flow_submodules = [
+                submodule
+                for _, submodule, _ in get_control_flow_submodules(current_graph_module)
+            ]
+            queue.extend(control_flow_submodules)
         return PassResult(graph_module, True)

exir/tests/TARGETS

@@ -204,6 +204,8 @@ python_unittest(
         ":lib",
         ":models",
         "//caffe2:torch",
+        "//caffe2/functorch:functorch_src",
+        "//executorch/exir:graph_module",
         "//executorch/exir:lib",
         "//executorch/exir:memory",
         "//executorch/exir:memory_planning",

exir/tests/test_passes.py

@@ -20,6 +20,7 @@
 from executorch.exir.dialects._ops import bind_pattern_to_op, ops, ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.emit import emit_program
+from executorch.exir.graph_module import get_control_flow_submodules
 from executorch.exir.pass_base import ExportPass, PassResult
 from executorch.exir.pass_manager import PassManager
 from executorch.exir.passes import (
@@ -45,6 +46,7 @@
 from executorch.exir.tests.common import register_additional_test_aten_ops
 from executorch.exir.tests.control_flow_models import FTCondDeadCode, FTMapBasic
 from executorch.exir.tests.models import MLP, Mul
+from functorch.experimental import control_flow
 
 from torch import nn
 from torch.fx import GraphModule, subgraph_rewriter
@@ -826,6 +828,70 @@ def test_debug_handle_generator_pass(self) -> None:
         for node in graph_module.graph.nodes:
             self.assertIn("debug_handle", node.meta)
 
+    def test_debug_handle_generator_pass_with_control_flow(self) -> None:
+        def true_nested(y: torch.Tensor) -> torch.Tensor:
+            y = y + y
+            y = torch.mm(y, y)
+            return y
+
+        def false_nested(y: torch.Tensor) -> torch.Tensor:
+            return torch.mm(y, y)
+
+        def true_fn(x: torch.Tensor, pred2: torch.Tensor) -> torch.Tensor:
+            z = control_flow.cond(pred2, true_nested, false_nested, [x])
+            return x + z
+
+        def false_fn(x: torch.Tensor, _) -> torch.Tensor:
+            return x.cos()
+
+        def map_fn(
+            x: torch.Tensor, pred1: torch.Tensor, pred2: torch.Tensor, y: torch.Tensor
+        ) -> torch.Tensor:
+            x = x.cos()
+            y = control_flow.cond(pred1, true_fn, false_fn, [y, pred2])
+            x = x + y
+            return x.sin()
+
+        def f(
+            xs: torch.Tensor, pred1: torch.Tensor, pred2: torch.Tensor, y: torch.Tensor
+        ) -> torch.Tensor:
+            y = torch.mm(y, y)
+            return control_flow.map(map_fn, xs, pred1, pred2, y)
+
+        inputs = (
+            torch.ones(2, 2),
+            torch.tensor([False]),
+            torch.tensor([False]),
+            torch.ones(2, 2),
+        )
+
+        graph_module = exir.capture(
+            f,
+            inputs,
+            exir.CaptureConfig(),
+        ).exported_program.graph_module
+
+        def check_debug_handle_metadata(graph_module: torch.fx.GraphModule) -> None:
+            queue = [graph_module]
+            while queue:
+                current_graph_module = queue.pop(0)
+                for node in current_graph_module.graph.nodes:
+                    self.assertIn("debug_handle", node.meta)
+                control_flow_submodules = [
+                    submodule
+                    for _, submodule, _ in get_control_flow_submodules(
+                        current_graph_module
+                    )
+                ]
+                queue.extend(control_flow_submodules)
+
+        DebugHandleGeneratorPass()(graph_module)
+        check_debug_handle_metadata(graph_module)
+
+        # Check debug handle still preserved after ScalarToTensorPass
+        ScalarToTensorPass()(graph_module)
+        check_debug_handle_metadata(graph_module)
+
     def test_symint_conversion(self) -> None:
         def f(x: torch.Tensor) -> torch.Tensor:
             return x + x.shape[0]