Add should_run() + fast-copy infrastructure with targeted_ops annotations (pytorch#18497)

Summary:
Pull Request resolved: pytorch#18497

Adds infrastructure for skipping and fast-copying unchanged nodes during
ExportPass execution, then annotates ~60 ARM backend passes to use it.

## Changes

### 1. should_run() hook on ExportPass / ArmPass
Subclasses that declare a `targeted_ops` class attribute (a set of op
overloads) can be skipped entirely when the graph contains none of their
target ops. ArmPass provides a default implementation via inheritance.

### 2. Fast-copy for cold nodes
When a pass declares `targeted_ops`, nodes whose ops are NOT in the set
are copied into the new graph via `graph.node_copy()` instead of full
FakeTensor dispatch. Per-node cost drops from ~0.4 ms to ~0.02 ms (~20x).

Includes a safety guard: nodes without `val` metadata (e.g. nodes
inserted by `call()` overrides before `super().call()`) fall back to
full dispatch instead of propagating None.

### 3. FakeTensor cache extension
Context manager `_extend_faketensor_cache_builtins()` temporarily extends
the FakeTensor dispatch cache to cover ExecuTorch op namespaces
(quantized_decomposed, tosa, dim_order_ops, cortex_m). Avoids redundant
re-dispatches for non-builtin ops across 50+ passes.

### 4. __init_subclass__ auto-discovery on ArmPass
Subclasses with existing `_TARGET_OPS`, `_supported_ops`, or
`_EDGE_OPS`/`_ATEN_OPS` attributes get `targeted_ops` populated
automatically at class definition time — no manual annotation needed.

### 5. targeted_ops annotations on ~60 ARM passes
Each annotation is a one-liner declaring the ops the pass checks in
`call_operator()`. Combined with should_run() and fast-copy, this
achieves the measured speedup below.

## Benchmark

Model: small CNN feature extractor (~50K params, 9 conv layers with
LayerNorm, targeting Ethos-U55 via the ARM/TOSA lowering pipeline).
Graph: ~1200 nodes, 146 ExportPass invocations.

  lower() before:  186 s
  lower() after:   100 s
  Passes skipped:  53 of 146
  Delta:           -86 s  (-46 %)
Adds should_run() hook to ExportPass that subclasses can override to skip
execution when a pass has no work to do. ArmPass implements a default that
checks a targeted_ops class attribute against the graph's call_function nodes.

Also adds:
- _fast_copy_node path in ExportInterpreter.run_node that uses graph.node_copy
  instead of full FakeTensor dispatch for cold nodes in passes that declare
  targeted_ops. Per-node cost drops from ~0.4ms to ~0.02ms.
- _extend_faketensor_cache_builtins context manager that extends FakeTensor
  dispatch cache to cover ExecuTorch ops (quantized_decomposed, tosa, etc.)
- __init_subclass__ on ArmPass for auto-discovery of targeted_ops from
  existing _TARGET_OPS, _supported_ops, _EDGE_OPS/_ATEN_OPS attributes
- targeted_ops annotations on ~60 ARM pass subclasses

Measured on SleepNet featurizer (U55 lowering):
  lower():  185s -> 96s  = -89s (-48%)

Differential Revision: D97528110

Author: Andrew Pullin

backends/arm/_passes/arm_pass.py

@@ -20,6 +20,25 @@
 class ArmPass(ExportPass):
     """Base class for Arm passes."""
 
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        if getattr(cls, "targeted_ops", None) is not None:
+            return
+        # Only auto-discover targeted_ops for passes that use the standard
+        # call_operator() pattern. Passes that override call() use _TARGET_OPS
+        # for their own graph manipulation logic, not as a fast-copy declaration.
+        if "call" in cls.__dict__:
+            return
+        for attr in ("_TARGET_OPS", "_supported_ops"):
+            ops = getattr(cls, attr, None)
+            if ops:
+                cls.targeted_ops = set(ops) if not isinstance(ops, set) else ops  # type: ignore[attr-defined]
+                return
+        edge = getattr(cls, "_EDGE_OPS", None)
+        aten = getattr(cls, "_ATEN_OPS", None)
+        if edge or aten:
+            cls.targeted_ops = {*(edge or ()), *(aten or ())}  # type: ignore[attr-defined]
+
     def __init__(self, tfa_pass: bool = False, *args, **kwargs) -> None:
         super().__init__(*args, **kwargs)
         self.submodule_depth = 0
@@ -78,6 +97,34 @@ def get_name(pass_) -> str:
                 f"Cannot get name for pass: {pass_}. It must be an instance of ExportPass or have a __name__ attribute."
             )
 
+    def should_run(self, graph_module: GraphModule) -> bool:
+        """Skip this pass if the graph contains none of its targeted ops.
+
+        Subclasses that define a ``targeted_ops`` class attribute (a set of
+        op overloads) get this check for free via inheritance.  Passes
+        without ``targeted_ops`` always run (the default).
+
+        Recursively checks control flow submodules (cond/while_loop) so
+        passes are not incorrectly skipped when targeted ops are nested.
+
+        """
+        targeted = getattr(self, "targeted_ops", None)
+        if targeted is None:
+            return True
+
+        from executorch.exir.graph_module import get_control_flow_submodules
+
+        def _has_targeted_op(gm: GraphModule) -> bool:
+            for node in gm.graph.nodes:
+                if node.op == "call_function" and node.target in targeted:
+                    return True
+            for _, submod, _ in get_control_flow_submodules(gm):
+                if _has_targeted_op(submod):
+                    return True
+            return False
+
+        return _has_targeted_op(graph_module)
+
     def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False):
         if not updated:
             return super().call_operator(op, args, kwargs, meta)

backends/arm/_passes/cast_to_int32_pass.py

@@ -6,9 +6,7 @@
 from typing import Set, Type
 
 import torch
-
 from executorch.backends.arm._passes.arm_pass import ArmPass
-
 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

backends/arm/_passes/conv1d_unsqueeze_pass.py

@@ -9,10 +9,8 @@
 from typing import Set, Type
 
 from executorch.backends.arm._passes import ArmPass
-
 from executorch.backends.arm._passes.rewrite_conv_pass import RewriteConvPass
 from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass
-
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -35,6 +33,8 @@ class Conv1dUnsqueezePass(ArmPass):
         SizeAdjustInputPass,
     }
 
+    targeted_ops = {exir_ops.edge.aten.convolution.default}
+
     def call_operator(self, op, args, kwargs, meta):
         if op != exir_ops.edge.aten.convolution.default:
             return super().call_operator(op, args, kwargs, meta)

backends/arm/_passes/convert_expand_copy_to_repeat.py

@@ -8,7 +8,6 @@
 from typing import cast, Set, Type
 
 import torch
-
 from executorch.backends.arm._passes.arm_pass import ArmPass
 from executorch.backends.arm._passes.unsqueeze_before_repeat_pass import (
     UnsqueezeBeforeRepeatPass,
@@ -58,6 +57,8 @@ class ConvertExpandCopyToRepeatPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = {UnsqueezeBeforeRepeatPass}
 
+    targeted_ops = {exir_ops.edge.aten.expand_copy.default}
+
     expand_copy = exir_ops.edge.aten.expand_copy.default
     repeat = exir_ops.edge.aten.repeat.default
 

backends/arm/_passes/convert_full_like_to_full_pass.py

@@ -9,7 +9,6 @@
 from executorch.backends.arm._passes.fuse_constant_ops_pass import (
     ComputeConstantOpsAOTPass,
 )
-
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -36,6 +35,8 @@ class ConvertFullLikeToFullPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = {ComputeConstantOpsAOTPass}
 
+    targeted_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,

backends/arm/_passes/convert_permute_singleton_to_view_pass.py

@@ -7,10 +7,8 @@
 from typing import Sequence, Set, Tuple, Type
 
 from executorch.backends.arm._passes.arm_pass import ArmPass
-
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
-
 from torch._ops import OpOverload
 
 
@@ -35,6 +33,8 @@ class ConvertPermuteSingletonToViewPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
+    targeted_ops = set(_PERMUTE_TARGETS)
+
     def call_operator(self, op, args, kwargs, meta):
         if op not in _PERMUTE_TARGETS:
             return super().call_operator(op, args, kwargs, meta)

backends/arm/_passes/convert_split_to_slice.py

@@ -21,6 +21,11 @@ class ConvertSplitToSlicePass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
+    targeted_ops = {
+        exir_ops.edge.aten.split_with_sizes_copy.default,
+        exir_ops.edge.aten.split_copy.Tensor,
+    }
+
     split_ops = (
         exir_ops.edge.aten.split_with_sizes_copy.default,
         exir_ops.edge.aten.split_copy.Tensor,

backends/arm/_passes/convert_squeezes_to_view.py

@@ -24,6 +24,11 @@ class ConvertSqueezesToViewPass(ArmPass):
 
     _passes_required_after: Set[Type[ExportPass]] = {FuseViewCopyTransformPass}
 
+    targeted_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,

backends/arm/_passes/convert_to_clamp_pass.py

@@ -6,11 +6,9 @@
 from typing import Set, Tuple, Type
 
 from executorch.backends.arm._passes import ArmPass
-
 from executorch.backends.arm._passes.fold_qdq_with_annotated_qparams_pass import (
     QuantizeClampArgumentsPass,
 )
-
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass
 
@@ -32,6 +30,8 @@ def get_clamp_params(op, args) -> Tuple[float | None, float | None]:
 class ConvertToClampPass(ArmPass):
     _passes_required_after: Set[Type[ExportPass]] = {QuantizeClampArgumentsPass}
 
+    targeted_ops = edge_operators
+
     def call_operator(self, op, args, kwargs, meta):
         if op not in edge_operators or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta)

backends/arm/_passes/decompose_acosh_pass.py

@@ -37,6 +37,8 @@ class DecomposeAcoshPass(ArmPass):
         MatchArgDtypePass,
     }
 
+    targeted_ops = {edge_acosh_op}
+
     def call_operator(self, op, args, kwargs, meta, updated=False):
 
         if op is not edge_acosh_op: