Arm backend: Don't run pass call_operator in submodules (#15736)

When a call_operator pass traces a graph with a submodule, it will use
call_operator to trace the nodes in the submodule, and the pass will
affect them. This means
1) Passes will operator on submodules both when their containing modules
run them and when the submodule itself is handled. 2) This doesn't
happen for call passes, so passes won't run in the intended order. All
call_operator passes will run before call passes.

To get around this, we detect submodules in ArmPass's call_submodule,
and use the default call_operator function when tracing them.

cc @freddan80 @per @zingo @oscarandersson8218 @digantdesai

Signed-off-by: Erik Lundell <[email protected]>

Author: Erik-Lundell

backends/arm/_passes/arm_pass.py

@@ -6,14 +6,20 @@
 
 import traceback
 from abc import abstractmethod
-from typing import List, Optional, Set, Type
+from typing import Any, List, Optional, Set, Type
 
 from executorch.exir.pass_base import ExportPass, NodeMetadata
+from torch.fx import GraphModule
+from torch.fx.passes.infra.pass_base import PassResult
 
 
 class ArmPass(ExportPass):
     """Base class for Arm passes"""
 
+    def __init__(self) -> None:
+        super().__init__()
+        self.submodule_depth = 0
+
     @property
     @abstractmethod
     def _passes_required_after(self) -> Set[Type[ExportPass]]:
@@ -56,3 +62,19 @@ def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False)
         old_stack_trace = new_meta.get("stack_trace", "")
         new_meta["stack_trace"] = f"{old_stack_trace}\n{traceback.format_stack()[-2]}"
         return super().call_operator(op, args, kwargs, NodeMetadata(new_meta))
+
+    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

backends/arm/test/misc/test_call_operator_submodule.py

@@ -0,0 +1,72 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any
+
+import torch
+
+from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager
+from executorch.backends.arm.tosa.specification import TosaSpecification
+from torch.fx import GraphModule
+from torch.fx.passes.infra.pass_base import PassResult
+
+
+class _DepthRecordingPass(ArmPass):
+    _passes_required_after = set()
+
+    def __init__(self, initial_graph_module):
+        super().__init__()
+        self.depths: list[int] = []
+        self.initial_submodule = initial_graph_module
+        self.submodule = None
+        self.num_submodules_called = 0
+
+    def call_operator(self, op, args, kwargs, meta, updated: bool = False):
+        """Should only be called from the top-level graph module."""
+        self.depths.append(self.submodule_depth)
+        assert self.submodule == self.initial_submodule
+        return super().call_operator(op, args, kwargs, meta, updated)
+
+    def call_submodule(
+        self, graph_module: GraphModule, inputs: tuple[Any, ...]
+    ) -> PassResult:
+        """Should be called for all three graph_modules: top-level, if, and else."""
+        self.submodule = graph_module
+        self.num_submodules_called += 1
+        return super().call_submodule(graph_module, inputs)
+
+
+class _CondModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        def _true_branch(arg: torch.Tensor) -> torch.Tensor:
+            return arg + 1
+
+        def _false_branch(arg: torch.Tensor) -> torch.Tensor:
+            return arg - 1
+
+        predicate = x.sum() > 0
+        return torch.cond(predicate, _true_branch, _false_branch, [x])
+
+
+def test_call_operator_runs_once_for_cond_submodules() -> None:
+    module = _CondModule()
+    example_inputs = (torch.randn(2, 3),)
+    exported = torch.export.export(module, example_inputs)
+    graph_module = exported.graph_module
+
+    recording_pass = _DepthRecordingPass(graph_module)
+    pass_manager = ArmPassManager(TosaSpecification.create_from_string("TOSA-1.00+FP"))
+    pass_manager.add_pass(recording_pass)
+    pass_manager._transform(graph_module)
+
+    assert recording_pass.num_submodules_called == 3
+    assert recording_pass.depths, "call_operator was never invoked"
+    assert (
+        max(recording_pass.depths) == 1
+    ), "call_operator was invoked with larger than one submodule depth."
+    assert (
+        min(recording_pass.depths) == 1
+    ), "call_operator was invoked with zero submodule depth."