Arm backend: Avoid running passes with no matching target ops (pytorch#19839)

Add ArmPass.should_run_pass() as a reusable early-exit hook before
  call() starts the normal ExportPass retracing path. The default hook
  returns true, preserving existing behavior for ArmPass subclasses.

  Introduce ArmOpTargetedPass for passes that only transform a known
  set of operator targets. It implements should_run_pass() by scanning
  the current graph and nested GraphModules for matching target
  operators. If no matching target operator is found, the pass returns
  an unmodified PassResult.

  For passes that already gate transformations with
  allowed_to_transform(), allow the target pre-scan to apply the same
  check before deciding whether the pass needs to run. This avoids
  running TFA passes when all matching target nodes are marked as
  disallowed.

  The should_run_pass() hook and ArmOpTargetedPass pre-scan avoid
  rebuilding graphs for decomposition and rewrite passes that cannot
  affect the current graph. The speedup is most visible on large models.

  Single-run paired benchmarks on Arm backend model tests
  across FP32, INT, VGF no-quant, and VGF quant variants:

  | Model       | E2E avg | Pass-manager avg |
  |-------------|--------:|-----------------:|
  | T5-small    | +30.5%  | +47.5%           |
  | DeepLabV3   | +12.9%  | +49.8%           |
  | Wav2Letter  | +16.9%  | +51.2%           |
  | InceptionV3 | +22.2%  | +46.5%           |
  | MobileNetV2 | +22.2%  | +52.5%           |
  | MobileNetV3 | +29.9%  | +54.6%           |

  Model rows are unweighted averages over successful variants.
  Unweighted average across 23 successful model/target variants:
  E2E speedup: +22.4%
  Pass-manager speedup: +50.5%

Change-Id: Iaa09638473a1d6d1e2ce98f5a0e3fc3a14378143


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

Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>
Co-authored-by: Erik Lundell <erik.lundell@arm.com>

Author: YufengShi-dudu

backends/arm/_passes/__init__.py

@@ -5,7 +5,7 @@
 
 
 from . import arm_pass_utils  # noqa
-from .arm_pass import ArmPass  # noqa  # usort: skip
+from .arm_pass import ArmOpTargetedPass, ArmPass  # noqa  # usort: skip
 from .accumulate_index_put_pass import AccumulateIndexPutPass  # noqa
 from .broadcast_args_pass import BroadcastArgsPass  # noqa
 from .canonicalize_gather_pass import CanonicalizeGatherPass  # noqa

backends/arm/_passes/accumulate_index_put_pass.py

@@ -6,7 +6,7 @@
 
 import torch
 
-from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.backends.arm._passes.decompose_index_tensor_to_gather_pass import (
     DecomposeIndexTensorToGatherPass,
 )
@@ -32,7 +32,7 @@ def get_ops(op):
     raise RuntimeError(f"Can't get index_put decomposition for op {op}")
 
 
-class AccumulateIndexPutPass(ArmPass):
+class AccumulateIndexPutPass(ArmOpTargetedPass):
     """This pass adjusts the values arg when the accumulate arg is set to true
     for the index_put op.
     """
@@ -41,9 +41,11 @@ class AccumulateIndexPutPass(ArmPass):
         DecomposeIndexTensorToGatherPass,
         RewriteIndexPutPass,
     }
+    target_ops = aten_ops + edge_ops
+    check_allowed_to_transform = True
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in (aten_ops + edge_ops) or not self.allowed_to_transform(meta):
+        if op not in self.target_ops or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta)
 
         source, indices, values = args[:3]

backends/arm/_passes/arm_pass.py

@@ -7,14 +7,15 @@
 import copy
 import traceback
 from abc import abstractmethod
+from collections.abc import Collection
 from typing import Any, List, Optional, Set, Type
 
 import torch
 from executorch.backends.arm.constants import DISALLOW_TFA_META_KEY
 from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue
-from torch.fx import GraphModule
+from torch.fx import GraphModule, Node
 from torch.fx.passes.infra.pass_base import PassResult
 from torch.utils import _pytree as pytree
 
@@ -191,3 +192,99 @@ def call_scalar(self, value: int | float, meta: NodeMetadata | dict[str, Any]):
             meta=meta,
             updated=True,
         )
+
+    def should_run_pass(self, graph_module: GraphModule) -> bool:
+        """Return whether this pass should run on the graph module.
+
+        Subclasses can override this to cheaply skip the pass before
+        ``call()`` starts the normal ``ExportPass`` retracing path.
+
+        Args:
+            graph_module (GraphModule): The graph module to inspect.
+
+        Returns:
+            bool: True when the pass should run.
+
+        """
+        return True
+
+    def __call__(self, graph_module: GraphModule) -> PassResult | None:
+        self.requires(graph_module)
+        if not self.should_run_pass(graph_module):
+            self.ensures(graph_module)
+            return PassResult(graph_module, False)
+        res = self.call(graph_module)
+        self.ensures(graph_module)
+        return res
+
+
+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.
+
+    Set ``check_allowed_to_transform`` to ``True`` when the target pre-scan
+    should also apply ``allowed_to_transform()`` to matching target nodes. This
+    is useful for TFA passes whose ``call_operator()`` leaves disallowed target
+    nodes unchanged. If all matching targets are disallowed, the pass can
+    return before entering the normal ``ExportPass`` path.
+
+    """
+
+    target_ops: Collection[Any] = ()
+    check_allowed_to_transform = False
+
+    def has_target_node(self, graph_module: GraphModule) -> bool:
+        """Return whether the graph module tree contains a target node.
+
+        Args:
+            graph_module (GraphModule): The graph module tree to inspect.
+
+        Returns:
+            bool: True if a matching call_function node is present.
+
+        """
+        visited_graph_modules = set()
+
+        def target_node_can_trigger_pass(node: Node) -> bool:
+            if not self.check_allowed_to_transform:
+                return True
+            if self.allowed_to_transform(node.meta):
+                return True
+            return False
+
+        def graph_has_target(module: GraphModule) -> bool:
+            if id(module) in visited_graph_modules:
+                return False
+            visited_graph_modules.add(id(module))
+
+            for target in self.target_ops:
+                for node in module.graph.find_nodes(
+                    op="call_function",
+                    target=target,
+                    sort=False,
+                ):
+                    if target_node_can_trigger_pass(node):
+                        return True
+
+            return any(
+                isinstance(child, GraphModule) and graph_has_target(child)
+                for child in module.children()
+            )
+
+        return graph_has_target(graph_module)
+
+    def should_run_pass(self, graph_module: GraphModule) -> bool:
+        """Return whether this pass has a target node to transform.
+
+        Args:
+            graph_module (GraphModule): The graph module tree to inspect.
+
+        Returns:
+            bool: True when a matching target node is present.
+
+        """
+        return self.has_target_node(graph_module)

backends/arm/_passes/canonicalize_gather_pass.py

@@ -6,12 +6,12 @@
 from typing import Set, Type
 
 import torch
-from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
 
-class CanonicalizeGatherPass(ArmPass):
+class CanonicalizeGatherPass(ArmOpTargetedPass):
     """Canonicalize gather so it can be lowered to TOSA.GATHER via the backend
     dialect.
 
@@ -40,10 +40,10 @@ class CanonicalizeGatherPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
-    _TARGET_OPS = {exir_ops.edge.aten.gather.default}
+    target_ops = {exir_ops.edge.aten.gather.default}
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in self._TARGET_OPS:
+        if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta)
 
         # edge.aten.gather.default: (x, dim, index) with kw-only sparse_grad

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -8,7 +8,7 @@
 
 from typing import Set, Type
 
-from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes import ArmOpTargetedPass
 
 from executorch.backends.arm._passes.rewrite_conv_pass import RewriteConvPass
 from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass
@@ -17,7 +17,7 @@
 from executorch.exir.pass_base import ExportPass
 
 
-class Conv1dUnsqueezePass(ArmPass):
+class Conv1dUnsqueezePass(ArmOpTargetedPass):
     """This pass is used to change conv1d ops into conv2d since TOSA only
     supports 2d and 3d convolution.
 
@@ -34,9 +34,10 @@ class Conv1dUnsqueezePass(ArmPass):
         RewriteConvPass,
         SizeAdjustInputPass,
     }
+    target_ops = (exir_ops.edge.aten.convolution.default,)
 
     def call_operator(self, op, args, kwargs, meta):
-        if op != exir_ops.edge.aten.convolution.default:
+        if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta)
         stride = list(args[3])
         if len(stride) != 1:

backends/arm/_passes/convert_expand_copy_to_repeat.py

@@ -9,7 +9,7 @@
 
 import torch
 
-from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass
 from executorch.backends.arm._passes.unsqueeze_before_repeat_pass import (
     UnsqueezeBeforeRepeatPass,
 )
@@ -51,7 +51,7 @@ def calculate_multiples(args):
     return multiples, expanded_rank != len(input_shape)
 
 
-class ConvertExpandCopyToRepeatPass(ArmPass):
+class ConvertExpandCopyToRepeatPass(ArmOpTargetedPass):
     """Replace expand copy with repeat since it is a repeat that can only repeat
     singleton dimensions.
     """
@@ -60,9 +60,10 @@ class ConvertExpandCopyToRepeatPass(ArmPass):
 
     expand_copy = exir_ops.edge.aten.expand_copy.default
     repeat = exir_ops.edge.aten.repeat.default
+    target_ops = (expand_copy,)
 
     def call_operator(self, op, args, kwargs, meta):
-        if op != self.expand_copy:
+        if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta)
 
         multiples, changes_rank = calculate_multiples(args)

backends/arm/_passes/convert_full_like_to_full_pass.py

@@ -5,7 +5,7 @@
 
 from typing import Set, Type
 
-from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass
 from executorch.backends.arm._passes.fuse_constant_ops_pass import (
     ComputeConstantOpsAOTPass,
 )
@@ -14,7 +14,7 @@
 from executorch.exir.pass_base import ExportPass
 
 
-class ConvertFullLikeToFullPass(ArmPass):
+class ConvertFullLikeToFullPass(ArmOpTargetedPass):
     """Convert edge aten full_like to full.
 
     As per the full_like PyTorch documentation, `torch.full_like(input,
@@ -35,11 +35,10 @@ class ConvertFullLikeToFullPass(ArmPass):
     """
 
     _passes_required_after: Set[Type[ExportPass]] = {ComputeConstantOpsAOTPass}
+    target_ops = (exir_ops.edge.aten.full_like.default,)
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in [
-            exir_ops.edge.aten.full_like.default,
-        ]:
+        if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta)
 
         tensor = args[0].data

backends/arm/_passes/convert_permute_singleton_to_view_pass.py

@@ -6,7 +6,7 @@
 
 from typing import Sequence, Set, Tuple, Type
 
-from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass
 
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
@@ -20,7 +20,7 @@
 )
 
 
-class ConvertPermuteSingletonToViewPass(ArmPass):
+class ConvertPermuteSingletonToViewPass(ArmOpTargetedPass):
     """Replace permutations that only move singleton axes with a reshape.
 
     Examples:
@@ -34,9 +34,10 @@ class ConvertPermuteSingletonToViewPass(ArmPass):
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()
+    target_ops = _PERMUTE_TARGETS
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in _PERMUTE_TARGETS:
+        if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta)
 
         input_tensor = args[0].data

backends/arm/_passes/convert_squeezes_to_view.py

@@ -6,15 +6,15 @@
 
 from typing import Set, Type
 
-from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.backends.arm._passes.fuse_view_copy_transform_pass import (
     FuseViewCopyTransformPass,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
 
-class ConvertSqueezesToViewPass(ArmPass):
+class ConvertSqueezesToViewPass(ArmOpTargetedPass):
     """Replaces squeeze/unsqueeze operators with view.
 
     These are simply special cases of the view op, so removing them gives us
@@ -23,12 +23,13 @@ class ConvertSqueezesToViewPass(ArmPass):
     """
 
     _passes_required_after: Set[Type[ExportPass]] = {FuseViewCopyTransformPass}
+    target_ops = (
+        exir_ops.edge.aten.squeeze_copy.dims,
+        exir_ops.edge.aten.unsqueeze_copy.default,
+    )
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in [
-            exir_ops.edge.aten.squeeze_copy.dims,
-            exir_ops.edge.aten.unsqueeze_copy.default,
-        ]:
+        if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta)
 
         x = args[0]

backends/arm/_passes/convert_to_clamp_pass.py

@@ -1,11 +1,11 @@
-# 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.
 
 from typing import Set, Tuple, Type
 
-from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes import ArmOpTargetedPass
 
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
     QuantizeClampArgumentsPass,
@@ -29,11 +29,13 @@ def get_clamp_params(op, args) -> Tuple[float | None, float | None]:
         raise ValueError(f"Getting clamp parameters for op {op} is not implemented.")
 
 
-class ConvertToClampPass(ArmPass):
+class ConvertToClampPass(ArmOpTargetedPass):
     _passes_required_after: Set[Type[ExportPass]] = {QuantizeClampArgumentsPass}
+    target_ops = edge_operators
+    check_allowed_to_transform = True
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in edge_operators or not self.allowed_to_transform(meta):
+        if op not in self.target_ops or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta)
 
         return super().call_operator(