Arm backend: Avoid unnecessary lowering pass retraces (pytorch#21138)

Add cheap applicability checks before selected Arm passes enter the
ExportPass retracing path. Skip InsertRescalePass for profiles without
integer support and graphs without DQ nodes. Move symbolic shape and
adaptive pool passes to ArmOpTargetedPass, and pre-scan
RemoveGetItemPass producers.

These checks preserve target-present transformations while avoiding
retraces on non-applicable graphs.

Lowering times are medians of three fresh-process exports with no
warm-up. Regular models use TOSA FP with FP32; full Qwen uses BF16.

| Model             | Base (s) | Updated (s) | Speedup |
|-------------------|---------:|------------:|--------:|
| ResNet-18         |    3.633 |       3.519 |   3.16% |
| MobileNetV3       |   11.063 |      10.100 |   8.70% |
| DeiT-tiny         |   33.323 |      31.137 |   6.56% |
| Conformer         |   13.121 |      12.204 |   6.98% |
| Test Qwen text    |   11.582 |      10.962 |   5.35% |
| Test Qwen vision  |   11.979 |      11.388 |   4.93% |
| Full Qwen text    |  155.825 |     146.824 |   5.78% |
| Full Qwen vision  |  126.671 |     120.153 |   5.15% |

Change-Id: I3c539e023efb5d1e59a725e531c850a897758002


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

Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>

Author: YufengShi-dudu

backends/arm/_passes/decompose_adaptive_max_pool2d_pass.py

@@ -6,7 +6,7 @@
 from typing import Set, Type
 
 import torch
-from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.backends.arm.constants import NHWC_INVERSE_ORDER, NHWC_ORDER
 from executorch.backends.arm.tosa.dialect.ops.max_pool2d import (
     compute_max_pool2d_output_shape,
@@ -16,7 +16,7 @@
 from executorch.exir.pass_base import ExportPass, NodeMetadata
 
 
-class DecomposeAdaptiveMaxPool2dPass(ArmPass):
+class DecomposeAdaptiveMaxPool2dPass(ArmOpTargetedPass):
     """Decompose irregular TOSA MAX_POOL2D_ADAPTIVE into per-bin slices.
 
     For dynamic-shape cases where ``MAX_POOL2D_ADAPTIVE`` cannot directly map
@@ -27,6 +27,7 @@ class DecomposeAdaptiveMaxPool2dPass(ArmPass):
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()
+    target_ops = {exir_ops.backend.tosa.MAX_POOL2D_ADAPTIVE.default}
 
     @staticmethod
     def _is_static_dim(dim) -> bool:

backends/arm/_passes/insert_rescales_pass.py

@@ -19,6 +19,7 @@
 from executorch.backends.arm._passes.quant_args import QuantArgs
 from executorch.backends.arm.constants import DQ_OPS, Q_OPS
 from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
+from executorch.backends.arm.tosa.specification import get_context_spec
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.fx import GraphModule, Node
@@ -110,6 +111,14 @@ def fold_dq_q_to_rescale(self, node: Node, user: Node, graph_module: GraphModule
             user.replace_all_uses_with(rescale_node)
             graph_module.graph.erase_node(user)
 
+    def should_run_pass(self, graph_module: GraphModule) -> bool:
+        if not get_context_spec().support_integer():
+            return False
+        return any(
+            graph_module.graph.find_nodes(op="call_function", target=target, sort=False)
+            for target in DQ_OPS
+        )
+
     def call(self, graph_module: GraphModule) -> PassResult:
         modified = False
         for node in graph_module.graph.nodes:

backends/arm/_passes/remove_getitem_pass.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.
@@ -12,3 +12,14 @@
 
 class RemoveGetItemPass(ArmPass, remove_getitem_op.RemoveGetItemPass):
     _passes_required_after: Set[Type[ExportPass]] = set()
+    _target_names = {
+        "aten.max_pool2d_with_indices.default",
+        "aten.max.dim",
+    }
+
+    def should_run_pass(self, graph_module) -> bool:
+        return any(
+            node.op == "call_function"
+            and getattr(node.target, "__name__", None) in self._target_names
+            for node in graph_module.graph.nodes
+        )

backends/arm/_passes/symbolic_to_tosa_shape_pass.py

@@ -7,13 +7,14 @@
 from typing import Optional
 
 import torch
-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
 
 
-class SymbolicToTosaShapesPass(ArmPass):
+class SymbolicToTosaShapesPass(ArmOpTargetedPass):
 
     _passes_required_after = set()
+    target_ops = {torch.ops.aten.sym_size.int}
 
     def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False):
         if op == torch.ops.aten.sym_size.int:

backends/arm/test/passes/test_ioquantization_pass.py

@@ -123,6 +123,33 @@ def _build_dq_q_graph(
     return builder.get_graph_module()
 
 
+def test_insert_rescale_should_run_based_on_spec_and_dq():
+    """Check integer support and DQ presence before running the pass."""
+    graph_module = _build_dq_q_graph(
+        torch.randint(-128, 127, (1, 4), dtype=torch.int8),
+        torch.int8,
+        torch.int8,
+        dq_scale=0.5,
+        dq_zp=0,
+        q_scale=0.25,
+        q_zp=0,
+    )
+    insert_rescale = InsertRescalePass()
+
+    with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+FP")):
+        assert not insert_rescale.should_run_pass(graph_module)
+
+    with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")):
+        assert insert_rescale.should_run_pass(graph_module)
+
+    builder = GraphBuilder()
+    x = builder.placeholder("x", torch.rand(1, 4))
+    builder.output([x])
+    graph_without_dq = builder.get_graph_module()
+    with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")):
+        assert not insert_rescale.should_run_pass(graph_without_dq)
+
+
 def test_insert_rescale_tosa_INT_folds_uint8_input():
     graph_module = _build_dq_q_graph(
         torch.randint(0, 255, (1, 4), dtype=torch.uint8),