Arm backend: Migrate pass manager to exported program (#20025)

Summary:
- Use ExportedProgramPassManager as the Arm pass manager base
- Keep GraphModule-only transforms on the non-deprecated FX manager

Test:
- PYTHONPATH=src:. /Users/usazah01/src/executorch/env/bin/python -m
pytest -q -p no:rerunfailures
backends/arm/test/misc/test_call_operator_submodule.py
backends/arm/test/passes/test_arm_pass_manager_insertions.py
backends/arm/test/misc/test_pass_pipeline_config.py
backends/arm/test/misc/test_pass_required_order.py

cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Usamah Zaheer <usamah.zaheer@arm.com>

Author: usamahz

backends/arm/_passes/arm_pass_manager.py

@@ -7,8 +7,9 @@
 
 import logging
 from collections import defaultdict
-from collections.abc import Sequence
+from collections.abc import Callable, Sequence
 from dataclasses import dataclass, field
+from typing import Any, cast
 
 from executorch.backends.arm._passes import (
     AccumulateIndexPutPass,
@@ -167,12 +168,17 @@
 )
 
 from executorch.exir import ExportedProgram
-from executorch.exir.pass_base import ExportPass
-from executorch.exir.pass_manager import PassManager
+from executorch.exir._program_utils import _get_updated_graph_signature
+from executorch.exir.pass_base import (
+    ExportedProgramPassBase,
+    ExportedProgramPassResult,
+    ExportPass,
+)
+from executorch.exir.pass_manager import ExportedProgramPassManager
 from torch._export.utils import _get_shape_env_from_gm
 from torch.fx import GraphModule
 from torch.fx.passes.infra.pass_base import PassResult
-from torch.nn.modules import Module
+from torch.fx.passes.infra.pass_manager import PassManager as GraphModulePassManager
 
 logger = logging.getLogger(__name__)
 
@@ -188,6 +194,50 @@ class PassInsertions:
 _registered_pass_insertions: dict[type, PassInsertions] = {}
 
 
+def _graph_pass_name(graph_pass: Callable[[GraphModule], PassResult | None]) -> str:
+    if isinstance(graph_pass, ExportPass):
+        return ArmPass.get_name(graph_pass)
+    if hasattr(graph_pass, "__name__"):
+        return graph_pass.__name__
+    return type(graph_pass).__name__
+
+
+class _ExportedProgramGraphPassAdapter(ExportedProgramPassBase):
+    def __init__(self, graph_pass: Callable[[GraphModule], PassResult | None]) -> None:
+        self.graph_pass = graph_pass
+
+    def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult:
+        graph_pass = cast(Any, self.graph_pass)
+        pass_exported_program = getattr(graph_pass, "exported_program", None)
+        if pass_exported_program is not None:
+            # ExportedProgramPassManager works on a shallow copy; Arm graph
+            # passes that store an ExportedProgram must update that copy.
+            graph_pass.exported_program = exported_program
+
+        try:
+            result = self.graph_pass(exported_program.graph_module)
+        finally:
+            if pass_exported_program is not None:
+                graph_pass.exported_program = pass_exported_program
+
+        if result is None:
+            raise TypeError(
+                f"The result of pass {_graph_pass_name(self.graph_pass)} should be type PassResult."
+            )
+
+        if result.modified:
+            result.graph_module.recompile()
+            exported_program._graph_module = result.graph_module
+            exported_program._graph_signature = _get_updated_graph_signature(
+                exported_program.graph_signature,
+                result.graph_module,
+            )
+            # Arm graph passes do not change symbolic shape constraints, and
+            # metadata-only fake modes may differ after propagation.
+
+        return ExportedProgramPassResult(exported_program, result.modified)
+
+
 def register_pass_insertions_before(
     target_pass_type: type, passes: list[ExportPass]
 ) -> None:
@@ -211,7 +261,7 @@ def clear_registered_pass_insertions() -> None:
     _registered_pass_insertions.clear()
 
 
-class ArmPassManager(PassManager):
+class ArmPassManager(ExportedProgramPassManager):
     def __init__(self, compile_spec: ArmCompileSpec) -> None:
         self.compile_spec = compile_spec
         self.tosa_spec = compile_spec.tosa_spec
@@ -374,8 +424,39 @@ def _tosa_context(self, graph_module: GraphModule) -> TosaLoweringContext:
         shape_env = _get_shape_env_from_gm(graph_module)
         return TosaLoweringContext(self.tosa_spec, shape_env)
 
-    def _transform(self, graph_module: GraphModule):
-        return self(graph_module).graph_module
+    def _transform_graph_module(self, graph_module: GraphModule):
+        # TFA and control-flow submodule paths operate on bare GraphModules
+        # without a standalone ExportedProgram to keep in sync.
+        return GraphModulePassManager(self.passes)(graph_module).graph_module
+
+    def __call__(  # type: ignore[override]
+        self,
+        module: ExportedProgram | GraphModule,
+        override_verifiers: Any | None = None,
+    ) -> ExportedProgramPassResult | PassResult:
+        if isinstance(module, GraphModule):
+            if override_verifiers is not None:
+                raise ValueError("override_verifiers is only valid for ExportedProgram")
+            return GraphModulePassManager(self.passes)(module)
+        return super().__call__(module, override_verifiers)
+
+    def _transform(
+        self,
+        exported_program: ExportedProgram,
+        graph_module: GraphModule,
+    ) -> GraphModule:
+        if graph_module is exported_program.graph_module:
+            passes: list[
+                ExportedProgramPassBase | Callable[[GraphModule], PassResult | None]
+            ] = [_ExportedProgramGraphPassAdapter(p) for p in self.passes]
+            transformed_program = ExportedProgramPassManager(passes)(
+                exported_program
+            ).exported_program
+            exported_program._graph_module = transformed_program.graph_module
+            exported_program._graph_signature = transformed_program.graph_signature
+            exported_program._range_constraints = transformed_program.range_constraints
+            return exported_program.graph_module
+        return self._transform_graph_module(graph_module)
 
     def add_pass(self, pipeline_pass):
         if type(pipeline_pass) in self._skip_pass_types:
@@ -558,7 +639,7 @@ def _tosa_pipeline(
         self._apply_pass_insertions()
 
         self.validate_constraints_mandatory()
-        return self._transform(graph_module)
+        return self._transform(exported_program, graph_module)
 
     def transform_to_backend_pipeline(
         self, exported_program: ExportedProgram, graph_module: GraphModule
@@ -663,21 +744,4 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 ]
             )
 
-            return self._transform(graph_module)
-
-    def __call__(self, module: Module) -> PassResult:
-        try:
-            return super().__call__(module)
-        except Exception as e:
-            first_exception = e.__cause__ or e.__context__ or e
-            import re
-
-            message = e.args[0]
-            m = re.search(r"An error occurred when running the '([^']+)' pass", message)
-            if m:
-                pass_name = m.group(1)
-                first_exception.args = (
-                    f"{pass_name}: {first_exception.args[0]}",
-                    *first_exception.args[1:],
-                )
-            raise first_exception
+            return self._transform_graph_module(graph_module)

backends/arm/test/misc/test_call_operator_submodule.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-2026 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.
@@ -60,7 +60,7 @@ def test_call_operator_runs_once_for_cond_submodules_tosa_FP() -> None:
     recording_pass = _DepthRecordingPass(graph_module)
     pass_manager = ArmPassManager(TosaCompileSpec("TOSA-1.00+FP"))
     pass_manager.add_pass(recording_pass)
-    pass_manager._transform(graph_module)
+    pass_manager._transform_graph_module(graph_module)
 
     assert recording_pass.num_submodules_called == 3
     assert recording_pass.depths, "call_operator was never invoked"

backends/arm/test/passes/test_arm_op_targeted_pass.py

@@ -4,7 +4,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import operator
-from typing import Set, Type
+from typing import cast, Set, Type
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass
@@ -45,7 +45,7 @@ def create_test_pass_manager() -> ArmPassManager:
 def run_single_pass(graph_module: GraphModule, test_pass: ExportPass) -> PassResult:
     pass_manager = create_test_pass_manager()
     pass_manager.add_pass(test_pass)
-    return pass_manager(graph_module)
+    return cast(PassResult, pass_manager(graph_module))
 
 
 class DummyTargetedPass(ArmOpTargetedPass):

exir/passes/__init__.py

@@ -32,7 +32,7 @@
     to_scratch_op,
 )
 from executorch.exir.pass_base import ExportPass
-from executorch.exir.pass_manager import PassManager, PassType
+from executorch.exir.pass_manager import ExportedProgramPassManager, PassType
 from executorch.exir.passes.const_prop_pass import ConstPropPass
 from executorch.exir.passes.debug_handle_generator_pass import DebugHandleGeneratorPass
 
@@ -498,25 +498,27 @@ def dead_code_elimination_pass(graph_module: torch.fx.GraphModule) -> PassResult
 
 # Passes to convert a graph module from ATen to Edge IR
 
-base_pre_op_replace_passes: List[Callable[[torch.nn.Module], PassResult]] = PassManager(
-    passes=[
-        # ReplaceSymSizeOpPass need to be run before other passes which inherits
-        # from ExportPass. ExportPass can not handle OpOverloadPacket in its
-        # call_function method. The ReplaceSymSizeOpPass pass converts sym size
-        # ops from OpOverloadPacket to OpOverload.
-        ReplaceSymSizeOpPass(),
-        NormalizeTransposePass(),
-        ReplaceBrokenOpsWithFunctionalOpsPass(),
-        ScalarToTensorPass(),
-        SymToTensorPass(),
-        RemoveNoopPass(),
-        PruneEmptyTensorsPass(),
-        RemoveToCopyPass(),
-    ]
-).passes
+base_pre_op_replace_passes: List[Callable[[torch.nn.Module], PassResult]] = (
+    ExportedProgramPassManager(
+        passes=[
+            # ReplaceSymSizeOpPass need to be run before other passes which inherits
+            # from ExportPass. ExportPass can not handle OpOverloadPacket in its
+            # call_function method. The ReplaceSymSizeOpPass pass converts sym size
+            # ops from OpOverloadPacket to OpOverload.
+            ReplaceSymSizeOpPass(),
+            NormalizeTransposePass(),
+            ReplaceBrokenOpsWithFunctionalOpsPass(),
+            ScalarToTensorPass(),
+            SymToTensorPass(),
+            RemoveNoopPass(),
+            PruneEmptyTensorsPass(),
+            RemoveToCopyPass(),
+        ]
+    ).passes
+)
 
 base_post_op_replace_passes: List[Callable[[torch.nn.Module], PassResult]] = (
-    PassManager(
+    ExportedProgramPassManager(
         passes=[
             dead_code_elimination_pass,
             DebugHandleGeneratorPass(),