Arm backend: Skip irrelevant lowering passes (#20659)

Add custom should_run_pass() pre-scans for Arm passes whose full
ExportPass retracing path is only useful for specific graph patterns.
The pre-scan keeps existing transforms unchanged when a relevant node is
present, but returns an unmodified PassResult when the graph cannot be
affected by the pass.

This avoids rebuilding graphs for placeholder, rank, dtype, buffer,
layout, no-op, and decomposition passes that would otherwise scan the
whole graph in call() and make no changes.

| Model       | Improvement |
|-------------|------------:|
| DeepLabV3   | +10.9%      |
| Wav2Letter  | +5.0%       |
| InceptionV3 | +3.2%       |
| MobileNetV2 | +12.8%      |
| MobileNetV3 | +12.1%      |

Signed-off-by: Sebastian Larsson <sebastian.larsson@arm.com>

Author: Sebastian-Larsson

backends/arm/_passes/broadcast_args_pass.py

@@ -28,20 +28,35 @@ class BroadcastArgsPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
-    targeted_ops = {
+    target_ops = {
         exir_ops.edge.aten.add.Tensor,
         exir_ops.edge.aten.sub.Tensor,
         # mul is indirectly targeting div as div is decompsed to reciprocal + mul
         exir_ops.edge.aten.mul.Tensor,
     }
 
-    def call(self, graph_module: GraphModule) -> PassResult:
+    def should_run_pass(self, graph_module: GraphModule) -> bool:
         tosa_spec = get_context_spec()
         if not tosa_spec.is_U55_subset:
-            return PassResult(graph_module, False)
+            return False
+        for node in graph_module.graph.nodes:
+            if node.op != "call_function" or node.target not in self.target_ops:
+                continue
+            output_shape = get_first_fake_tensor(node).shape
+            broadcasts = 0
+            for arg in node.args:
+                if not isinstance(arg, Node):
+                    continue
+                if get_first_fake_tensor(arg).shape != output_shape:
+                    broadcasts += 1
+                if broadcasts > 1:
+                    return True
+        return False
+
+    def call(self, graph_module: GraphModule) -> PassResult:
         modified = False
         for node in graph_module.graph.nodes:
-            if node.op != "call_function" or node.target not in self.targeted_ops:
+            if node.op != "call_function" or node.target not in self.target_ops:
                 continue
 
             output_shape = get_first_fake_tensor(node).shape

backends/arm/_passes/cast_int64_pass.py

@@ -25,6 +25,19 @@ def __init__(self, exported_program: ExportedProgram, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.exported_program = exported_program
 
+    def should_run_pass(self, graph_module: torch.fx.GraphModule) -> bool:
+        for node in graph_module.graph.nodes:
+            if len(node.users) == 0:
+                continue
+            fake_tensor = node.meta.get("val")
+            if not isinstance(fake_tensor, torch._subclasses.fake_tensor.FakeTensor):
+                continue
+            if fake_tensor.dtype == torch.int64 and is_buffer(
+                self.exported_program, node
+            ):
+                return True
+        return False
+
     def _assert_within_int32(self, tensor: torch.Tensor, node: torch.fx.Node):
         if torch.min(tensor) < torch.iinfo(torch.int32).min:
             raise RuntimeError(

backends/arm/_passes/cast_to_int32_pass.py

@@ -7,33 +7,33 @@
 
 import torch
 
-from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass
 
 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 executorch.exir.pass_base import ExportPass
 
 
-class CastToInt32Pass(ArmPass):
+class CastToInt32Pass(ArmOpTargetedPass):
     """Casts the input to int32 if it is not already and casts back the output
     to the original input dtype.
     """
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
-    targeted_ops = {
+    target_ops = {
         exir_ops.edge.aten.bitwise_left_shift.Tensor,
         exir_ops.edge.aten.bitwise_right_shift.Tensor,
     }
 
-    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+    def should_run_pass(self, graph_module: torch.fx.GraphModule) -> bool:
         tosa_spec = get_context_spec()
         if not tosa_spec.is_U55_subset:
-            return PassResult(graph_module, False)
-        return super().call(graph_module)
+            return False
+        return super().should_run_pass(graph_module)
 
     def call_operator(self, op, args, kwargs, meta):
-        if op not in self.targeted_ops:
+        if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta)
 
         new_args: list = []

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -15,6 +15,7 @@
 
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
+from torch.fx import GraphModule
 
 
 class Conv1dUnsqueezePass(ArmOpTargetedPass):
@@ -36,6 +37,17 @@ class Conv1dUnsqueezePass(ArmOpTargetedPass):
     }
     target_ops = (exir_ops.edge.aten.convolution.default,)
 
+    def should_run_pass(self, graph_module: GraphModule) -> bool:
+        for node in graph_module.graph.nodes:
+            if node.op != "call_function" or node.target not in self.target_ops:
+                continue
+            if len(node.args) > 3 and len(node.args[3]) == 1:
+                return True
+        return any(
+            isinstance(child, GraphModule) and self.should_run_pass(child)
+            for child in graph_module.children()
+        )
+
     def call_operator(self, op, args, kwargs, meta):
         if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta)

backends/arm/_passes/convert_split_to_slice.py

@@ -7,7 +7,7 @@
 from typing import Set, Type
 
 import torch.fx
-from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.backends.arm._passes.arm_pass_utils import (
     create_node,
     get_first_fake_tensor,
@@ -16,7 +16,7 @@
 from executorch.exir.pass_base import ExportPass, PassResult
 
 
-class ConvertSplitToSlicePass(ArmPass):
+class ConvertSplitToSlicePass(ArmOpTargetedPass):
     """Replace a split operation with many slice operations."""
 
     _passes_required_after: Set[Type[ExportPass]] = set()
@@ -25,6 +25,7 @@ class ConvertSplitToSlicePass(ArmPass):
         exir_ops.edge.aten.split_with_sizes_copy.default,
         exir_ops.edge.aten.split_copy.Tensor,
     )
+    target_ops = split_ops
     slice = exir_ops.edge.aten.slice_copy.Tensor
 
     def call(self, graph_module: torch.fx.GraphModule):

backends/arm/_passes/decompose_add_sub_alpha_pass.py

@@ -12,6 +12,7 @@
 from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
+from torch.fx import GraphModule
 
 
 _ADD_OPS = (
@@ -61,6 +62,19 @@ class DecomposeAddSubAlphaPass(ArmOpTargetedPass):
     _passes_required_after: Set[Type[ExportPass]] = set()
     target_ops = _ADD_OPS + _SUB_OPS
 
+    def should_run_pass(self, graph_module: GraphModule) -> bool:
+        for node in graph_module.graph.nodes:
+            if (
+                node.op == "call_function"
+                and node.target in self.target_ops
+                and _should_decompose(node.kwargs.get("alpha", 1))
+            ):
+                return True
+        return any(
+            isinstance(child, GraphModule) and self.should_run_pass(child)
+            for child in graph_module.children()
+        )
+
     def call_operator(self, op, args, kwargs, meta, updated: bool | None = False):
         if op not in self.target_ops:
             return super().call_operator(op, args, kwargs, meta, updated)

backends/arm/_passes/decompose_batch_norm_no_stats.py

@@ -8,7 +8,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._passes.arm_pass_utils import create_node
 from executorch.backends.arm._passes.fuse_constant_ops_pass import (
     ComputeConstantOpsAOTPass,
@@ -19,7 +19,7 @@
 from executorch.exir.pass_base import ExportPass, PassResult
 
 
-class DecomposeBatchNormNoStatsPass(ArmPass):
+class DecomposeBatchNormNoStatsPass(ArmOpTargetedPass):
     """
     Decompose BatchNorm2d(track_running_stats=False) (aten._native_batch_norm_legit_no_training)
     into a sequence of elementwise operations:
@@ -42,21 +42,21 @@ class DecomposeBatchNormNoStatsPass(ArmPass):
         ComputeConstantOpsAOTPass,
         InsertTableOpsPass,
     }
+    target_ops = (
+        exir_ops.edge.aten._native_batch_norm_legit.no_stats,
+        exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
+        torch.ops.aten._native_batch_norm_legit_no_training.default,
+        torch.ops.aten.batch_norm.default,
+        torch.ops.aten.native_batch_norm.default,
+    )
+    check_allowed_to_transform = True
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
-        bn_ops = (
-            exir_ops.edge.aten._native_batch_norm_legit.no_stats,
-            exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
-            torch.ops.aten._native_batch_norm_legit_no_training.default,
-            torch.ops.aten.batch_norm.default,
-            torch.ops.aten.native_batch_norm.default,
-        )
-
         modified = False
         for node in graph_module.graph.nodes:
             if (
                 node.op != "call_function"
-                or node.target not in bn_ops
+                or node.target not in self.target_ops
                 or not self.allowed_to_transform(node.meta)
             ):
                 continue

backends/arm/_passes/decompose_dynamic_full_pass.py

@@ -6,15 +6,15 @@
 from typing import Any, Set, Type
 
 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,
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
 
-class DecomposeDynamicFullPass(ArmPass):
+class DecomposeDynamicFullPass(ArmOpTargetedPass):
     """Rewrite dynamic-shape `full` into scalar `full` plus `repeat`."""
 
     _passes_required_after: Set[Type[ExportPass]] = {UnsqueezeBeforeRepeatPass}
@@ -23,6 +23,7 @@ class DecomposeDynamicFullPass(ArmPass):
         torch.ops.aten.full.default,
         exir_ops.edge.aten.full.default,
     }
+    target_ops = full_targets
     repeat = exir_ops.edge.aten.repeat.default
 
     @staticmethod

backends/arm/_passes/decompose_layernorm_pass.py

@@ -8,7 +8,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._passes.arm_pass_utils import create_node, insert_scalar
 from executorch.backends.arm._passes.decompose_meandim_pass import DecomposeMeanDimPass
 from executorch.backends.arm._passes.decompose_var_pass import DecomposeVarPass
@@ -46,7 +46,7 @@ def get_layer_norm_decomposition(op) -> tuple:
     raise RuntimeError(f"Can't get layer_norm composition for op {op}")
 
 
-class DecomposeLayerNormPass(ArmPass):
+class DecomposeLayerNormPass(ArmOpTargetedPass):
     """
     layernorm is defined as: ((x - E[x]) / sqrt(Var[x] + eps)) * weights + bias
     Decompose layernorm(x, normalized_shape, weights, bias, eps) to a sequence of:
@@ -73,6 +73,8 @@ class DecomposeLayerNormPass(ArmPass):
         exir_ops.edge.aten.native_layer_norm.default,
         torch.ops.aten.layer_norm.default,
     }
+    target_ops = _TARGET_OPS
+    check_allowed_to_transform = True
 
     def call(self, graph_module: torch.fx.GraphModule):
         modified = False

backends/arm/_passes/decompose_linear_pass.py

@@ -8,7 +8,7 @@
 
 import torch
 
-from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.backends.arm._passes.arm_pass_utils import (
     create_node,
     get_first_fake_tensor,
@@ -18,7 +18,7 @@
 from executorch.exir.pass_base import ExportPass, PassResult
 
 
-class DecomposeLinearPass(ArmPass):
+class DecomposeLinearPass(ArmOpTargetedPass):
     """This pass decomposes linear into a Conv2D with view operations.
 
     Example:
@@ -31,6 +31,7 @@ class DecomposeLinearPass(ArmPass):
     """
 
     _passes_required_after: Set[Type[ExportPass]] = {InsertRescaleInt32Pass}
+    target_ops = (exir_ops.edge.aten.linear.default,)
 
     def call(self, graph_module):
         modified = False