Update on "Use caching allocator for runner (#15730)"

Summary:

We observed that on iOS it improves perf by 6% because SDPA op does temp allocations.

No significant difference on android though.
ghstack-source-id: 328001114
exported-using-ghexport

Reviewed By: navsud, derekdixu

Differential Revision: D86120038




[ghstack-poisoned]

Author: kimishpatel

.ci/scripts/test_cortex_m_e2e.sh

@@ -1,5 +1,6 @@
 #!/usr/bin/env bash
 # Copyright (c) Meta Platforms, Inc. and affiliates.
+# Copyright 2026 Arm Limited and/or its affiliates.
 # All rights reserved.
 #
 # This source code is licensed under the BSD-style license found in the
@@ -18,7 +19,7 @@ mkdir -p "./cortex_m_e2e/${MODEL}"
 WORK_DIR=$(realpath "./cortex_m_e2e/${MODEL}")
 
 echo "=== Exporting ${MODEL} with cortex-m55+int8 ==="
-python -m examples.arm.aot_arm_compiler \
+python -m backends.arm.scripts.aot_arm_compiler \
     -m "${MODEL}" \
     --target=cortex-m55+int8 \
     --quantize \

backends/arm/README.md

@@ -106,6 +106,13 @@ Setup:
 ./examples/arm/setup.sh --disable-ethos-u-deps --enable-mlsdk-deps
 ```
 
+This is the default setup path and installs the MLSDK components from pip.
+Developers who need local source builds can use:
+
+```
+./backends/arm/scripts/setup-mlsdk-from-source.sh
+```
+
 The current flow lowers to TOSA and converts to VGF for use in external projects,
 so the `executor_runner` is not typically used here.
 
@@ -155,7 +162,7 @@ scp -P 2222 arm_test/cmake-out/executor_runner root@127.0.0.1:/tmp/
 Create a PTE file:
 
 ```
-python3 -m examples.arm.aot_arm_compiler \
+python3 -m backends.arm.scripts.aot_arm_compiler \
   --model_name examples/arm/example_modules/add.py \
   --delegate \
   --quantize \

backends/arm/_passes/__init__.py

@@ -143,6 +143,7 @@
 from .replace_scalar_with_tensor_pass import (  # noqa
     ReplaceScalarWithTensorByProfilePass,
 )
+from .rewrite_avg_pool2d_pass import RewriteAvgPool2dPass  # noqa
 from .rewrite_bool_bitwise_to_logical_pass import (  # noqa
     RewriteBoolBitwiseToLogicalPass,
 )
@@ -157,6 +158,7 @@
 from .rewrite_inplace_arithmetic_pass import RewriteInplaceArithmeticPass  # noqa
 from .rewrite_le_lt_to_ge_gt_pass import RewriteLeLtToGeGtPass  # noqa
 from .rewrite_matmul import RewriteMatmulPass  # noqa
+from .rewrite_max_pool2d_pass import RewriteMaxPool2dPass  # noqa
 from .rewrite_pad import RewritePadPass  # noqa
 from .rewrite_slice import RewriteSlicePass  # noqa
 from .rewrite_upsample import RewriteUpsamplePass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -10,7 +10,6 @@
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 
-import executorch.backends.arm.tosa.dialect  # noqa: unused
 from executorch.backends.arm._passes import (
     AccumulateIndexPutPass,
     AnnotateOutputDimOrderPass,
@@ -126,6 +125,7 @@
     RemoveNoopPass,
     ReplaceInfAndLimitValuesPass,
     ReplaceScalarWithTensorByProfilePass,
+    RewriteAvgPool2dPass,
     RewriteBoolBitwiseToLogicalPass,
     RewriteBoolToFp32CastViaInt8Pass,
     RewriteConvPass,
@@ -134,6 +134,7 @@
     RewriteInplaceArithmeticPass,
     RewriteLeLtToGeGtPass,
     RewriteMatmulPass,
+    RewriteMaxPool2dPass,
     RewritePadPass,
     RewriteSlicePass,
     RewriteUpsamplePass,
@@ -143,7 +144,6 @@
     UnsqueezeBeforeRepeatPass,
     UnsqueezeScalarPlaceholdersPass,
 )
-
 from executorch.backends.arm._passes.arm_pass import ArmPass
 from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec
 from executorch.backends.arm.common.pipeline_config import (
@@ -175,6 +175,32 @@ class PassInsertions:
     after_passes: list = field(default_factory=list)
 
 
+_registered_pass_insertions: dict[type, PassInsertions] = {}
+
+
+def register_pass_insertions_before(
+    target_pass_type: type, passes: list[ExportPass]
+) -> None:
+    """Register passes to be inserted before a target pass for all pipelines."""
+    if target_pass_type not in _registered_pass_insertions:
+        _registered_pass_insertions[target_pass_type] = PassInsertions()
+    _registered_pass_insertions[target_pass_type].before_passes.extend(passes)
+
+
+def register_pass_insertions_after(
+    target_pass_type: type, passes: list[ExportPass]
+) -> None:
+    """Register passes to be inserted after a target pass for all pipelines."""
+    if target_pass_type not in _registered_pass_insertions:
+        _registered_pass_insertions[target_pass_type] = PassInsertions()
+    _registered_pass_insertions[target_pass_type].after_passes.extend(passes)
+
+
+def clear_registered_pass_insertions() -> None:
+    """Clear all globally registered pass insertions."""
+    _registered_pass_insertions.clear()
+
+
 class ArmPassManager(PassManager):
     def __init__(self, compile_spec: ArmCompileSpec) -> None:
         self.compile_spec = compile_spec
@@ -323,13 +349,17 @@ 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.
+        Subclasses can override this to call insert_passes_before/after.
 
         Args:
             exported_program: The exported program being transformed
 
         """
-        pass
+        for pass_type, insertions in _registered_pass_insertions.items():
+            if insertions.before_passes:
+                self.insert_passes_before(pass_type, list(insertions.before_passes))
+            if insertions.after_passes:
+                self.insert_passes_after(pass_type, list(insertions.after_passes))
 
     def add_passes(self, passes: Sequence[ExportPass | None]):
         for p in passes:
@@ -432,6 +462,8 @@ def _tosa_pipeline(
                 DecomposeSliceScatterPass(),
                 AccumulateIndexPutPass(),
                 DecomposeIndexTensorToGatherPass(),
+                DecomposeAdaptiveAvgPool2dPass(),
+                DecomposeAvgPool2dPass(),
                 Conv1dUnsqueezePass(),
             ]
         )
@@ -468,17 +500,16 @@ def _tosa_pipeline(
                 DecomposeSoftmaxPass(),
                 ConvertMinMaxPass(),
                 DecomposeAnyPass(),
-                DecomposeAdaptiveAvgPool2dPass(),
-                DecomposeAvgPool2dPass(),
                 DecorateFp32toInt32CastingPass(),
-                ComputeConstantOpsAOTPass(exported_program),
-                FuseConstantArgsPass(exported_program),
                 ConvertExpandCopyToRepeatPass(),
                 UnsqueezeBeforeRepeatPass(),
                 DecomposeCumsumPass(exported_program),
                 DecomposeAsStridedCopyPass(),
                 DecomposeMaxPool2dPass(),
                 SizeAdjustInputPass(),
+                RewriteAvgPool2dPass(),
+                ComputeConstantOpsAOTPass(exported_program),
+                FuseConstantArgsPass(exported_program),
                 DecomposeSelectPass(),
                 ConvertSqueezesToViewPass(),
                 CastToInt32Pass(),
@@ -496,6 +527,7 @@ def _tosa_pipeline(
         self.add_passes(
             [
                 RewriteUpsamplePass(),
+                RewriteMaxPool2dPass(),
                 RewriteConvPass(exported_program),
                 RewriteMatmulPass(),
                 RewritePadPass(),
@@ -573,6 +605,8 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                     DecomposeLayerNormPass(tfa_pass=True),
                     DecomposeVarPass(tfa_pass=True),
                     DecomposeMeanDimPass(graph_module, self.tosa_spec, tfa_pass=True),
+                    DecomposeAdaptiveAvgPool2dPass(tfa_pass=True),
+                    DecomposeAvgPool2dPass(tfa_pass=True),
                 ]
             )
 
@@ -598,8 +632,6 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                     DecomposeDivPass(tfa_pass=True),
                     DecomposeLinalgVectorNormPass(tfa_pass=True),
                     DecomposeSqrtPass(tfa_pass=True),
-                    DecomposeAdaptiveAvgPool2dPass(tfa_pass=True),
-                    DecomposeAvgPool2dPass(tfa_pass=True),
                     DecomposeSoftmaxPass(
                         tfa_pass=True,
                     ),