Implement scheme for custom arm_pass_manager insertions (#16396)

Differential Revision: D87120925

Pull Request resolved: #16396

Author: 3l1

backends/arm/_passes/TARGETS

@@ -1,13 +1,55 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 runtime.python_library(
-    name = "passes",
-    srcs = glob(["*.py"]),
+    name = "core",
+    srcs = [
+        "arm_pass.py",
+        "arm_pass_utils.py",
+        "quant_args.py",
+    ],
     deps = [
         "//executorch/backends/arm:common",
         "//executorch/backends/arm:constants",
         "//executorch/backends/arm/tosa:utils",
         "//executorch/backends/arm/tosa/dialect:lib",
+        "//executorch/backends/transforms:utils",
+        "//executorch/exir:lib",
+    ],
+)
+
+runtime.python_library(
+    name = "arm_pass_manager_base",
+    srcs = ["arm_pass_manager.py"],
+    deps = [
+        ":core",
+        "//executorch/backends/arm:common",
+        "//executorch/backends/arm/tosa:utils",
+        "//executorch/backends/arm/tosa/dialect:lib",
+        "//executorch/exir:lib",
+    ],
+)
+
+runtime.python_library(
+    name = "arm_pass_manager_fb",
+    srcs = [],
+    deps = [
+        ":arm_pass_manager_base",
+        # @oss-disable[end= ]: "//executorch/backends/arm/_passes/fb:fb",
+    ],
+)
+
+runtime.python_library(
+    name = "passes",
+    srcs = glob(["*.py"], exclude = [
+        "arm_pass.py",
+        "arm_pass_utils.py",
+        "quant_args.py",
+    ]),
+    deps = [
+        ":core",
+        ":arm_pass_manager_base" if runtime.is_oss else ":arm_pass_manager_fb",
+        "//executorch/backends/arm/tosa:utils",
+        "//executorch/backends/arm/tosa/dialect:lib",
         "//executorch/backends/transforms:fuse_view_copy",
         "//executorch/backends/transforms:remove_getitem_op",
         "//executorch/backends/transforms:replace_scalar_with_tensor",

backends/arm/_passes/__init__.py

@@ -158,4 +158,13 @@
 from .control_flow_const_inline import (  # noqa  # usort: skip
     ControlFlowConstInlinePass,
 )
-from .arm_pass_manager import ArmPassManager  # noqa  # usort: skip
+
+# Import all subpackages to allow extensions to patch classes
+import importlib  # noqa: E402
+import pkgutil  # noqa: E402
+
+for _, _modname, _ispkg in pkgutil.iter_modules(__path__, __name__ + "."):
+    if _ispkg:
+        importlib.import_module(_modname)
+
+from .arm_pass_manager import ArmPassManager  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -8,6 +8,7 @@
 import logging
 from collections import defaultdict
 from collections.abc import Sequence
+from dataclasses import dataclass, field
 
 import executorch.backends.arm.tosa.dialect  # noqa: unused
 from executorch.backends.arm._passes import (
@@ -157,11 +158,21 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class PassInsertions:
+    """Holds lists of passes to be inserted before and after a target pass."""
+
+    before_passes: list = field(default_factory=list)
+    after_passes: list = field(default_factory=list)
+
+
 class ArmPassManager(PassManager):
     def __init__(self, compile_spec: ArmCompileSpec) -> None:
         self.compile_spec = compile_spec
         self.tosa_spec = compile_spec.tosa_spec
         self._skip_pass_types: tuple[type, ...] = ()
+        self._pass_insertions: dict[type, PassInsertions] = {}
+        self._insertions_applied = False
         super().__init__()
         self.configure_skip_passes()
 
@@ -223,6 +234,98 @@ def validate_constraints_mandatory(self):
 
             raise RuntimeError(error_msg)
 
+    def insert_passes_before(
+        self, target_pass_type: type, passes: list[ExportPass]
+    ) -> None:
+        """Register passes to be inserted before instances of target_pass_type.
+        Insertions are deferred and applied via _apply_pass_insertions().
+
+        Args:
+            target_pass_type: The pass class to insert before (e.g., InsertTableOpsPass)
+            passes: List of pass instances to insert
+
+        """
+        self._pass_insertions.setdefault(
+            target_pass_type, PassInsertions()
+        ).before_passes.extend(passes)
+
+    def insert_passes_after(
+        self, target_pass_type: type, passes: list[ExportPass]
+    ) -> None:
+        """Register passes to be inserted after instances of target_pass_type.
+        Insertions are deferred and applied via _apply_pass_insertions().
+
+        Args:
+            target_pass_type: The pass class to insert after
+            passes: List of pass instances to insert
+
+        """
+        self._pass_insertions.setdefault(
+            target_pass_type, PassInsertions()
+        ).after_passes.extend(passes)
+
+    def _apply_pass_insertions(self) -> None:
+        """Apply all registered pass insertions to the collected passes.
+
+        Called ONCE after all add_passes() calls are complete, before execution.
+
+        Raises:
+            ValueError: If any registered target pass type is not found in the pipeline.
+
+        """
+        if self._insertions_applied or not self._pass_insertions:
+            return
+
+        # Fail fast if any target pass type is missing from the pipeline
+        existing_pass_types = {type(p) for p in self.passes}
+        for target_type in self._pass_insertions:
+            if target_type not in existing_pass_types:
+                available = [type(p).__name__ for p in self.passes]
+                raise ValueError(
+                    f"Target pass {target_type.__name__} not found in the pass "
+                    f"pipeline. Available passes: {available}"
+                )
+
+        # Build new pass list with insertions applied
+        new_passes = []
+        for pass_obj in self.passes:
+            pass_type = type(pass_obj)
+
+            # Insert passes BEFORE this pass
+            if pass_type in self._pass_insertions:
+                insertions = self._pass_insertions[pass_type]
+                for before_pass in insertions.before_passes:
+                    # Check if we should skip this inserted pass
+                    if type(before_pass) not in self._skip_pass_types:
+                        new_passes.append(before_pass)
+
+            # Add the original pass
+            new_passes.append(pass_obj)
+
+            # Insert passes AFTER this pass
+            if pass_type in self._pass_insertions:
+                insertions = self._pass_insertions[pass_type]
+                for after_pass in insertions.after_passes:
+                    # Check if we should skip this inserted pass
+                    if type(after_pass) not in self._skip_pass_types:
+                        new_passes.append(after_pass)
+
+        # Replace the passes list
+        self.passes = new_passes
+        self._insertions_applied = True
+
+    def _configure_pass_insertions(self, exported_program: ExportedProgram) -> None:
+        """Hook for subclasses to configure pass insertions. Called at the START
+        of pipeline construction, before any passes are added.
+
+        Subclasses should override this to call insert_passes_before/after.
+
+        Args:
+            exported_program: The exported program being transformed
+
+        """
+        pass
+
     def add_passes(self, passes: Sequence[ExportPass | None]):
         for p in passes:
             if p is not None:
@@ -241,6 +344,9 @@ def add_pass(self, pipeline_pass):
     def _tosa_pipeline(
         self, exported_program: ExportedProgram, graph_module: GraphModule
     ) -> GraphModule:
+        # Allow subclasses to configure pass insertions before building pipeline
+        self._configure_pass_insertions(exported_program)
+
         # Preprocessing passes
         self.add_pass(AnnotateOutputDimOrderPass())
 
@@ -398,6 +504,9 @@ def _tosa_pipeline(
             ]
         )
 
+        # Apply all pass insertions once after all passes are collected
+        self._apply_pass_insertions()
+
         self.validate_constraints_mandatory()
         return self._transform(graph_module)
 

backends/arm/test/ops/test_rsqrt.py

@@ -146,7 +146,7 @@ def test_rsqrt_tosa_INT_a16w8(test_tensor: torch.Tensor):
 
 @common.parametrize("test_tensor", Rsqrt.test_parameters)
 @common.XfailIfNoCorstone300
-def test_rsqrt_16a8w_u55_INT16(test_tensor: torch.Tensor):
+def test_rsqrt_16a8w_u55_INT(test_tensor: torch.Tensor):
     """Test rsqrt operation with int16 I/O quantization for U55."""
     # Use wider tolerances for int16 I/O quantization on U55
     pipeline = EthosU55PipelineINT[input_t1](