Fast-copy cold nodes during ExportPass retracing (#18497)

backends/arm/_passes/arm_pass.py

@@ -1,3 +1,3 @@
 # Copyright 2025-2026 Arm Limited and/or its affiliates.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -24,6 +24,12 @@
 from torch.utils import _pytree as pytree
 
 
+_OPS_WITHOUT_QUANTIZED_FAKE_KERNEL = (
+    exir_ops.edge.aten.bmm.default,
+    exir_ops.edge.aten.leaky_relu.default,
+)
+
+
 class ArmPass(ExportPass):
     """Base class for Arm passes."""
 
@@ -57,6 +63,13 @@
         output_qparams = meta_dict.get("output_qparams", {})
         return bool(input_qparams) and bool(output_qparams)
 
+    def should_fast_copy_node(self, target: torch.fx.node.Target) -> bool:
+        # Metadata is unavailable here, so conservatively replay every op whose
+        # quantized path needs a custom fake result.
+        if any(target is op for op in _OPS_WITHOUT_QUANTIZED_FAKE_KERNEL):
+            return False
+        return super().should_fast_copy_node(target)
+
     @property
     @abstractmethod
     def _passes_required_after(self) -> Set[Type[ExportPass]]:
@@ -86,12 +99,8 @@
             )
 
     def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False):
-        ops_without_quantized_fake_kernel = {
-            exir_ops.edge.aten.bmm.default,
-            exir_ops.edge.aten.leaky_relu.default,
-        }
         if (
-            op in ops_without_quantized_fake_kernel
+            op in _OPS_WITHOUT_QUANTIZED_FAKE_KERNEL
             and isinstance(meta, NodeMetadata)
             and len(meta.data.get("input_qparams", {})) > 0
         ):
@@ -138,21 +147,32 @@
         self.tracer.set_metadata(res_proxy.node, res_data)
         return ProxyValue(res_data, res_proxy)
 
+    def should_run_on_nested_submodule(self, graph_module: GraphModule) -> bool:
+        """Return whether subclass rewrites should run in this nested submodule."""
+        has_custom_precheck = type(self).should_run_pass is not ArmPass.should_run_pass
+        return has_custom_precheck and self.should_run_pass(graph_module)
+
     def call_submodule(
         self, graph_module: GraphModule, inputs: tuple[Any, ...]
     ) -> PassResult:
         self.submodule_depth += 1
-        if self.submodule_depth == 1:
-            result = super().call_submodule(graph_module, inputs)
-        else:
-            # When we trace a submodule, we don't want to apply the calling pass.
-            # Temporarily replace call_operator to avoid this.
-            _call_operator_fn = self.call_operator
-            self.call_operator = super().call_operator  # type: ignore
-            result = super().call_submodule(graph_module, inputs)
-            self.call_operator = _call_operator_fn  # type: ignore
-        self.submodule_depth -= 1
-        return result
+        try:
+            if self.submodule_depth == 1:
+                self._top_level_call_operator = self.call_operator
+                return super().call_submodule(graph_module, inputs)
+
+            call_operator = self.call_operator
+            try:
+                if self.should_run_on_nested_submodule(graph_module):
+                    self.call_operator = self._top_level_call_operator  # type: ignore
+                else:
+                    # Nested submodules still need replay without subclass rewrites.
+                    self.call_operator = super().call_operator  # type: ignore
+                return super().call_submodule(graph_module, inputs)
+            finally:
+                self.call_operator = call_operator  # type: ignore
+        finally:
+            self.submodule_depth -= 1
 
     def call_shape_operator(
         self, op, args: tuple, kwargs: dict, meta: NodeMetadata, updated: bool = True
@@ -235,10 +255,15 @@
 class ArmOpTargetedPass(ArmPass):
     """Base class for passes that only transform selected operators.
 
-    Subclasses set ``target_ops`` to the call_function targets they can
-    transform. If the current graph and nested control-flow subgraphs do not
-    contain any target, the pass returns immediately without paying the default
-    ExportPass retracing cost.
+    Subclasses must set ``target_ops`` to the exhaustive set of call_function
+    targets their ``call_operator()`` can transform. This ARM-specific contract
+    drives both the target pre-scan and fast-copy eligibility; an empty
+    ``target_ops`` disables fast copy and skips the pass. Generic ``ExportPass``
+    subclasses instead use ``targeted_ops`` for explicit fast-copy opt-in.
+
+    If the current graph and nested control-flow subgraphs do not contain any
+    target, the pass returns immediately without paying the default ExportPass
+    retracing cost.
 
     Set ``check_allowed_to_transform`` to ``True`` when the target pre-scan
     should also apply ``allowed_to_transform()`` to matching target nodes. This
@@ -248,14 +273,28 @@
 
     """
 
+    enable_fast_copy = True
     target_ops: Collection[Any] = ()
     check_allowed_to_transform = False
 
-    def has_target_node(self, graph_module: GraphModule) -> bool:
+    def get_fast_copy_target_ops(self) -> Optional[tuple[Any, ...]]:
+        if not self.enable_fast_copy:
+            return None
+        targets = tuple(self.target_ops)
+        return targets if targets else None
+
+    def should_run_on_nested_submodule(self, graph_module: GraphModule) -> bool:
+        return self.has_target_node(graph_module, recursive=False)
+
+    def has_target_node(
+        self, graph_module: GraphModule, recursive: bool = True
+    ) -> bool:
         """Return whether the graph module tree contains a target node.
 
         Args:
             graph_module (GraphModule): The graph module tree to inspect.
+            recursive (bool): Whether to inspect nested child GraphModules.
+                When false, inspect only ``graph_module`` itself.
 
         Returns:
             bool: True if a matching call_function node is present.
@@ -284,7 +323,7 @@
                     if target_node_can_trigger_pass(node):
                         return True
 
-            return any(
+            return recursive and any(
                 isinstance(child, GraphModule) and graph_has_target(child)
                 for child in module.children()
             )

backends/arm/test/passes/test_arm_op_targeted_pass.py

@@ -15,6 +15,7 @@
 from executorch.exir.pass_base import ExportPass
 from torch.fx import Graph, GraphModule
 from torch.fx.passes.infra.pass_base import PassResult
+from torch.fx.passes.shape_prop import _extract_tensor_metadata
 
 
 TARGET_OP = torch.ops.aten.add.Tensor
@@ -50,15 +51,17 @@ def run_single_pass(graph_module: GraphModule, test_pass: ExportPass) -> PassRes
 
 class DummyTargetedPass(ArmOpTargetedPass):
     _passes_required_after: Set[Type[ExportPass]] = set()
-    target_ops = (TARGET_OP,)
+    target_ops: tuple[torch._ops.OpOverload, ...] = (TARGET_OP,)
     check_allowed_to_transform = True
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.call_operator_count = 0
+        self.called_ops = []
 
     def call_operator(self, op, args, kwargs, meta):
         self.call_operator_count += 1
+        self.called_ops.append(op)
         return super().call_operator(op, args, kwargs, meta)
 
 
@@ -80,6 +83,28 @@ def call(self, graph_module: GraphModule) -> PassResult:
         return PassResult(graph_module, True)
 
 
+class NestedCondModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def outer_true(arg: torch.Tensor) -> torch.Tensor:
+            def inner_true(inner_arg: torch.Tensor) -> torch.Tensor:
+                return inner_arg + inner_arg
+
+            def inner_false(inner_arg: torch.Tensor) -> torch.Tensor:
+                return inner_arg * inner_arg
+
+            return torch.cond(
+                arg.sum() > 1,
+                inner_true,
+                inner_false,
+                [arg],
+            )
+
+        def outer_false(arg: torch.Tensor) -> torch.Tensor:
+            return arg * arg
+
+        return torch.cond(x.sum() > 0, outer_true, outer_false, [x])
+
+
 class CondModule(torch.nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         def true_branch(arg: torch.Tensor) -> torch.Tensor:
@@ -91,6 +116,43 @@ def false_branch(arg: torch.Tensor) -> torch.Tensor:
         return torch.cond(x.sum() > 0, true_branch, false_branch, [x])
 
 
+def test_replays_targeted_op_and_fast_copies_cold_op() -> None:
+    graph = Graph()
+    lhs = graph.placeholder("lhs")
+    rhs = graph.placeholder("rhs")
+    lhs.meta["val"] = torch.randn(2, 3)
+    rhs.meta["val"] = torch.randn(2, 3)
+    cold = graph.call_function(torch.ops.aten.mul.Tensor, (lhs, rhs))
+    cold.meta["val"] = lhs.meta["val"] * rhs.meta["val"]
+    cold.meta["tensor_meta"] = _extract_tensor_metadata(cold.meta["val"])
+    targeted = graph.call_function(TARGET_OP, (cold, rhs))
+    targeted.meta["val"] = cold.meta["val"] + rhs.meta["val"]
+    targeted.meta["tensor_meta"] = _extract_tensor_metadata(targeted.meta["val"])
+    graph.output(targeted)
+    graph_module = GraphModule(torch.nn.Module(), graph)
+    targeted_pass = DummyTargetedPass()
+
+    result = run_single_pass(graph_module, targeted_pass)
+
+    assert result.modified
+    assert targeted_pass.call_operator_count == 1
+
+
+def test_empty_target_ops_disables_fast_copy_and_skips_pass() -> None:
+    class EmptyTargetedPass(DummyTargetedPass):
+        target_ops: tuple[torch._ops.OpOverload, ...] = ()
+
+    graph_module = create_graph_module(TARGET_OP)
+    targeted_pass = EmptyTargetedPass()
+
+    result = run_single_pass(graph_module, targeted_pass)
+
+    assert targeted_pass.get_fast_copy_target_ops() is None
+    assert result.graph_module is graph_module
+    assert not result.modified
+    assert targeted_pass.call_operator_count == 0
+
+
 def test_skips_when_target_is_absent() -> None:
     graph_module = create_graph_module()
     targeted_pass = DummyTargetedPass()
@@ -138,6 +200,18 @@ def test_runs_when_previous_pass_creates_target() -> None:
     assert targeted_pass.call_operator_count == 1
 
 
+def test_runs_when_target_is_present_in_deeply_nested_submodule() -> None:
+    exported_program = torch.export.export(NestedCondModule(), (torch.randn(2, 3),))
+    graph_module = exported_program.graph_module
+    targeted_pass = DummyTargetedPass()
+
+    result = run_single_pass(graph_module, targeted_pass)
+
+    assert result is not None
+    assert result.modified
+    assert TARGET_OP in targeted_pass.called_ops
+
+
 def test_runs_when_target_is_present_in_nested_submodule() -> None:
     exported_program = torch.export.export(CondModule(), (torch.randn(2, 3),))
     graph_module = exported_program.graph_module
@@ -147,4 +221,4 @@ def test_runs_when_target_is_present_in_nested_submodule() -> None:
 
     assert result is not None
     assert result.modified
-    assert targeted_pass.call_operator_count > 0
+    assert TARGET_OP in targeted_pass.called_ops