Add targeted_ops to core ExportPass passes for fast-copy optimization (#21702)

backends/arm/_passes/arm_pass.py

@@ -140,6 +140,8 @@ def call_submodule(
         self.submodule_depth += 1
         if self.submodule_depth == 1:
             result = super().call_submodule(graph_module, inputs)
+        elif self.should_run_pass(graph_module):
+            result = super().call_submodule(graph_module, inputs)
         else:
             # When we trace a submodule, we don't want to apply the calling pass.
             # Temporarily replace call_operator to avoid this.

exir/pass_base.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 # Copyright 2025-2026 Arm Limited and/or its affiliates.
@@ -143,6 +143,37 @@
     return None
 
 
+def _tensor_metadata_changed(original: Argument, new: Argument) -> bool:
+    original_leaves, original_spec = pytree.tree_flatten(original)
+    new_leaves, new_spec = pytree.tree_flatten(new)
+    if original_spec != new_spec:
+        return True
+
+    for original_leaf, new_leaf in zip(original_leaves, new_leaves):
+        original_is_tensor = isinstance(original_leaf, torch.Tensor)
+        new_is_tensor = isinstance(new_leaf, torch.Tensor)
+        if original_is_tensor != new_is_tensor:
+            return True
+        if not original_is_tensor:
+            continue
+
+        if (
+            original_leaf.shape != new_leaf.shape
+            or original_leaf.dtype != new_leaf.dtype
+            or original_leaf.layout != new_leaf.layout
+            or original_leaf.device != new_leaf.device
+            or original_leaf.requires_grad != new_leaf.requires_grad
+        ):
+            return True
+        if (
+            original_leaf.layout == torch.strided
+            and original_leaf.stride() != new_leaf.stride()
+        ):
+            return True
+
+    return False
+
+
 class NodeMetadata:
     def __init__(self, data: Dict[str, Any]) -> None:
         self.data: Dict[str, Any] = data.copy()
@@ -228,6 +259,10 @@
     pass
 
 
+class _FastCopyFallback(RuntimeError):
+    pass
+
+
 @dataclass(frozen=True)
 class ExportedProgramPassResult:
     exported_program: ExportedProgram
@@ -280,6 +315,22 @@
         """
 
 
+_FAST_COPY_UNSAFE_TARGETS: frozenset[Any] = frozenset(
+    {
+        torch.ops.aten.convolution,
+        torch.ops.aten.convolution.default,
+        torch.ops.aten.linear,
+        torch.ops.aten.linear.default,
+    }
+)
+
+
+def _is_fast_copy_unsafe_target(target: Any) -> bool:
+    if isinstance(target, EdgeOpOverload):
+        target = target._op
+    return target in _FAST_COPY_UNSAFE_TARGETS
+
+
 class _ExportPassBase(PassBase):
     """
     Interpreter-based pass class to help users maintain the IR spec while writing
@@ -413,12 +464,62 @@
 
             node.meta["tensor_meta"] = pytree.tree_map(make_tensor_meta, value)
 
+    # Types whose nodes are eligible for the fast-copy optimisation in
+    # ``run_node``.  Subclass interpreters (e.g. ``ExportPass``) extend
+    # this tuple to include dialect-specific overload types such as
+    # ``EdgeOpOverload``.
+    _OPERATOR_TARGET_TYPES: Tuple[type, ...] = (
+        torch._ops.OpOverload,
+        torch._ops.OpOverloadPacket,
+    )
+
     class ExportInterpreter(fx.Interpreter):
         def __init__(self, callback: "_ExportPassBase", gm: fx.GraphModule) -> None:
             super().__init__(gm)
             self.callback = callback
             self.node: torch.fx.Node = next(iter(gm.graph.nodes))
 
+            # --- fast-copy bookkeeping ---------------------------------
+            # When the owning pass declares ``targeted_ops``, cold nodes
+            # (those whose target is *not* in the set) can be copied into
+            # the new graph without an expensive FakeTensor dispatch.
+            targeted = getattr(callback, "targeted_ops", None)
+            if targeted is None:
+                targeted = getattr(callback, "target_ops", None)
+            if targeted is not None:
+                try:
+                    targeted_set = set(targeted)
+                except TypeError:
+                    targeted_set = None
+                self._targeted_ops: Optional[Set[Any]] = targeted_set
+            else:
+                self._targeted_ops: Optional[Set[Any]] = None
+
+            # Fast-copy relies on the existing ``n.meta["val"]`` being
+            # correct for cold nodes.  If the pass overrides ``call()``
+            # it may modify the graph (e.g. insert nodes with metadata
+            # copied from unrelated ops) before calling ``super().call()``,
+            # which would make cold-node metadata unreliable.  Disable the
+            # optimisation in that case.
+            call_overridden = type(callback).call is not _ExportPassBase.call
+            has_problematic_target = False
+            if self._targeted_ops:
+                for t in self._targeted_ops:
+                    if _is_fast_copy_unsafe_target(t):
+                        has_problematic_target = True
+                        break
+            self._fast_copy_enabled: bool = (
+                self._targeted_ops is not None
+                and not call_overridden
+                and not has_problematic_target
+            )
+
+            # Maps old-graph nodes to their new-graph equivalents so that
+            # ``_fast_copy_node`` can remap arguments (including get_attr
+            # nodes that are stored in ``self.env`` as raw tensors rather
+            # than ProxyValues).
+            self._node_remap: Dict[torch.fx.Node, torch.fx.Node] = {}
+
         def placeholder(  # pyre-fixme[14]
             self,
             target: str,
@@ -512,10 +613,153 @@
         ) -> None:
             raise ExportPassBaseError("call_method is not supported.")
 
+        # -- fast-copy helpers ------------------------------------------
+
+        def _preflight_fast_copy_inputs(
+            self,
+            n: torch.fx.Node,
+            tracer: "_ExportPassBase.ExportTracer",
+        ) -> Dict[torch.fx.Node, Tuple[Any, List[str]]]:
+            get_attr_values: Dict[torch.fx.Node, Tuple[Any, List[str]]] = {}
+            # Fallback must happen before copying nodes or creating module paths.
+            for old_node in n.all_input_nodes:
+                if old_node in self._node_remap:
+                    continue
+                pv = self.env.get(old_node)
+                if pv is not None and hasattr(pv, "proxy"):
+                    continue
+                if old_node.op != "get_attr":
+                    raise _FastCopyFallback
+
+                target_atoms = old_node.target.split(".")
+                root = tracer.root
+                for atom in target_atoms[:-1]:
+                    if not hasattr(root, atom):
+                        # The fresh tracer root receives this path after preflight.
+                        break
+                    child = getattr(root, atom)
+                    if not isinstance(child, torch.nn.Module):
+                        raise _FastCopyFallback
+                    root = child
+                get_attr_values[old_node] = (
+                    self.fetch_attr(old_node.target),
+                    target_atoms,
+                )
+            return get_attr_values
+
+        @staticmethod
+        def _ensure_get_attr_parent(
+            tracer: "_ExportPassBase.ExportTracer",
+            target_atoms: List[str],
+        ) -> torch.nn.Module:
+            root = tracer.root
+            for atom in target_atoms[:-1]:
+                if hasattr(root, atom):
+                    child = getattr(root, atom)
+                    if not isinstance(child, torch.nn.Module):
+                        raise _FastCopyFallback
+                else:
+                    child = torch.nn.Module()
+                    setattr(root, atom, child)
+                root = child
+            return root
+
+        def _fast_copy_arg(
+            self,
+            old_node: torch.fx.Node,
+            tracer: "_ExportPassBase.ExportTracer",
+            get_attr_values: Dict[torch.fx.Node, Tuple[Any, List[str]]],
+        ) -> torch.fx.Node:
+            new_node = self._node_remap.get(old_node)
+            if new_node is not None:
+                return new_node
+
+            proxy_value = self.env.get(old_node)
+            if proxy_value is not None and hasattr(proxy_value, "proxy"):
+                mapped = proxy_value.proxy.node
+                self._node_remap[old_node] = mapped
+                return mapped
+
+            if old_node.op != "get_attr":
+                raise _FastCopyFallback
+
+            attribute: Optional[Tuple[torch.nn.Module, str, Any]] = None
+            if old_node.op == "get_attr":
+                value, target_atoms = get_attr_values[old_node]
+                attribute = (
+                    self._ensure_get_attr_parent(tracer, target_atoms),
+                    target_atoms[-1],
+                    value,
+                )
+
+            copied = tracer.graph.node_copy(
+                old_node, lambda node: self._node_remap.get(node, node)
+            )
+            self._node_remap[old_node] = copied
+            if attribute is not None:
+                root, name, value = attribute
+                setattr(root, name, value)
+            return copied
+
+        def _fast_copy_node(self, n: torch.fx.Node) -> "ProxyValue":
+            tracer = self.callback.tracer
+            get_attr_values = self._preflight_fast_copy_inputs(n, tracer)
+
+            new_node = tracer.graph.node_copy(
+                n,
+                lambda old_node: self._fast_copy_arg(
+                    old_node, tracer, get_attr_values
+                ),
+            )
+
+            val = n.meta.get("val")
+            proxy = torch.fx.Proxy(new_node, tracer)
+            result = ProxyValue(val, proxy)
+            self._node_remap[n] = new_node
+            return result
+
         def run_node(self, n: torch.fx.Node) -> Argument:
             self.node = n
             self.callback.node_debug_str = n.format_node()
-            return super().run_node(n)
+
+            # Fast-copy path: skip the full interpreter dispatch for cold
+            # call_function nodes whose operator is not targeted by this
+            # pass.  This avoids the expensive FakeTensor re-dispatch and
+            # proxy reconstruction for nodes the pass will not modify.
+            if (
+                self._fast_copy_enabled
+                and n.op == "call_function"
+                and isinstance(n.target, self.callback._OPERATOR_TARGET_TYPES)
+                and n.target not in self._targeted_ops  # type: ignore[operator]
+                and n.meta.get("val") is not None
+            ):
+                try:
+                    return self._fast_copy_node(n)
+                except _FastCopyFallback:
+                    self._fast_copy_enabled = False
+
+            result = super().run_node(n)
+
+            # Record old→new node mapping for fast-copy arg remapping.
+            if self._fast_copy_enabled and isinstance(result, ProxyValue):
+                self._node_remap[n] = result.proxy.node
+
+            # After a hot node runs through full dispatch, verify that
+            # it did not change tensor metadata. If it did, downstream
+            # cold nodes' original ``val`` metadata would be stale, so
+            # we disable the fast-copy optimisation for the remainder
+            # of this interpreter walk.
+            if (
+                self._fast_copy_enabled
+                and n.op == "call_function"
+                and self._targeted_ops is not None
+                and n.target in self._targeted_ops
+                and isinstance(result, ProxyValue)
+            ):
+                if _tensor_metadata_changed(n.meta.get("val"), result.data):
+                    self._fast_copy_enabled = False
+
+            return result
 
     def __init__(self) -> None:
         self.interpreter = torch.fx.Interpreter(
@@ -768,13 +1012,17 @@
     def call_submodule(
         self, graph_module: fx.GraphModule, inputs: Tuple[Argument, ...]
     ) -> PassResult:
-        prev_tracer, self.tracer = self.tracer, self.ExportTracer(
-            self, graph_module.graph._codegen
+        prev_tracer, self.tracer = (
+            self.tracer,
+            self.ExportTracer(self, graph_module.graph._codegen),
         )
         self.tracer.fake_tensor_mode = prev_tracer.fake_tensor_mode
         interpreter = self.ExportInterpreter(self, graph_module)
-        prev_interpreter, self.interpreter = self.interpreter, torch.fx.Interpreter(
-            torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
+        prev_interpreter, self.interpreter = (
+            self.interpreter,
+            torch.fx.Interpreter(
+                torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())
+            ),
         )
         inputs_data = pytree.tree_map_only(ProxyValue, lambda x: x.data, inputs)
         with fx_traceback.preserve_node_meta():
@@ -824,6 +1072,14 @@
 
 
 class ExportPass(_ExportPassBase):
+    # Extend operator target types to include the Edge dialect overloads so
+    # that the fast-copy optimisation in ``run_node`` also covers Edge ops.
+    _OPERATOR_TARGET_TYPES: Tuple[type, ...] = (
+        torch._ops.OpOverload,
+        torch._ops.OpOverloadPacket,
+        EdgeOpOverload,
+    )
+
     class ExportTracer(_ExportPassBase.ExportTracer):
         def create_arg(self, a: Argument) -> torch.fx.Node:
             if isinstance(a, torch.nn.Module):

exir/passes/memory_format_ops_pass.py

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -29,6 +29,8 @@
     the aten op and the new edge dialect dim_order op.
     """
 
+    targeted_ops = DimOrderOpsMap.keys()
+
     def call_operator(self, op, args, kwargs, meta):
         if not (isinstance(op, EdgeOpOverload) and op in DimOrderOpsMap):
             return super().call_operator(
@@ -57,9 +59,9 @@
         elif isinstance(args[0], torch.fx.immutable_collections.immutable_list):
             ndim = len(args[0])
         else:
-            assert (
-                0
-            ), f"Expecting a Tensor, a ProxyValue, or a Sequence, but got {type(args[0])}"
+            assert 0, (
+                f"Expecting a Tensor, a ProxyValue, or a Sequence, but got {type(args[0])}"
+            )
 
         # Derive dim_order based on memory format
         dim_order: List[int]
@@ -96,6 +98,8 @@
     This pass is to revert the dim_order ops back to the memory format ops.
     """
 
+    targeted_ops = MemoryFormatOpsMap.keys()
+
     def call_operator(self, op, args, kwargs, meta):
         if not (isinstance(op, EdgeOpOverload) and op in MemoryFormatOpsMap):
             return super().call_operator(

exir/passes/normalize_transpose_pass.py

@@ -13,9 +13,12 @@ class NormalizeTransposePass(ExportPass):
     Even with functionalization on, we still get graph with
     torch.ops.aten.t.default op. Ideally we should fix functionalization.
     TODO: once we have that, we should remove this pass.
-    Check test_normalize_transpose_op in test_passes.py for more details
+    Check test_normalize_transpose_uses_targeted_fast_copy in
+    test_pass_infra.py for more details.
     """
 
+    targeted_ops = {torch.ops.aten.t.default}
+
     def call_operator(self, op, args, kwargs, meta):
         if op == torch.ops.aten.t.default:
             return super().call_operator(

exir/passes/remove_mixed_type_operators.py

@@ -14,6 +14,15 @@
 
 
 class RemoveMixedTypeOperators(ExportPass):
+    targeted_ops = {
+        torch.ops.aten.add.Tensor,
+        torch.ops.aten.mul.Tensor,
+        torch.ops.aten.sub.Tensor,
+        torch.ops.aten.div.Tensor,
+        torch.ops.aten.div.Tensor_mode,
+        torch.ops.aten.minimum.default,
+    }
+
     # pyre-ignore
     def call_operator(self, op, args, kwargs, meta: NodeMetadata):  # noqa: C901
         if len(args) <= 1: