From 984e6359916fb777a7f640c93e18f1489ee4b7f1 Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Sun, 6 Sep 2026 09:47:30 -0700 Subject: [PATCH] Fast-copy cold nodes during ExportPass retracing (#18497) Summary: Optimize `ExportPass` replay for passes that declare `target_ops` or `targeted_ops` by copying cold `call_function` nodes with `graph.node_copy` instead of redispatching them through FakeTensor. Old-to-new node remapping preserves dependencies and `get_attr` values. The fast path validates all inputs before mutating the destination graph or module tree, so unsupported remapping falls back without leaving partial state. An explicitly empty `targeted_ops` takes precedence over legacy `target_ops`. Fast-copy is disabled when `call()` is overridden, for exact convolution or linear targets, and after a hot node changes nested tensor metadata. Nested ARM control-flow submodules continue to use `ArmPass.should_run_pass()`. A/B benchmarking on CombinedControl U55 lowering, with each revision run twice, showed a 12.5% speedup. The synthetic U55 suite was within run-to-run noise. Earlier versions of this diff also contained the ARM pass-skipping implementation. That code was copied into https://github.com/pytorch/executorch/pull/19839 (D106781989) and landed under ARM authorship; this diff now contains the remaining fast-copy optimization. Differential Revision: D97528110 --- backends/arm/_passes/arm_pass.py | 15 +- exir/pass_base.py | 321 ++++++++++++++++++++++++++++++- exir/tests/test_pass_infra.py | 314 ++++++++++++++++++++++++++++++ 3 files changed, 642 insertions(+), 8 deletions(-) diff --git a/backends/arm/_passes/arm_pass.py b/backends/arm/_passes/arm_pass.py index 5c3541a7586..28ecc5362dc 100644 --- a/backends/arm/_passes/arm_pass.py +++ b/backends/arm/_passes/arm_pass.py @@ -57,6 +57,15 @@ def _is_quantized_meta(self, meta: NodeMetadata | dict[str, Any]) -> bool: output_qparams = meta_dict.get("output_qparams", {}) return bool(input_qparams) and bool(output_qparams) + def should_fast_copy_node(self, target: torch.fx.node.Target) -> bool: + ops_without_quantized_fake_kernel = { + exir_ops.edge.aten.bmm.default, + exir_ops.edge.aten.leaky_relu.default, + } + if any(target is op for op in ops_without_quantized_fake_kernel): + return False + return super().should_fast_copy_node(target) + @property @abstractmethod def _passes_required_after(self) -> Set[Type[ExportPass]]: @@ -142,11 +151,11 @@ def call_submodule( self, graph_module: GraphModule, inputs: tuple[Any, ...] ) -> PassResult: self.submodule_depth += 1 - if self.submodule_depth == 1: + if self.submodule_depth == 1 or 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. + # Nested submodules that do not need this pass still need normal replay. + # Temporarily replace call_operator to avoid applying subclass rewrites. _call_operator_fn = self.call_operator self.call_operator = super().call_operator # type: ignore result = super().call_submodule(graph_module, inputs) diff --git a/exir/pass_base.py b/exir/pass_base.py index 6071aae2be8..eba2c8722c7 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -103,6 +103,16 @@ class _SymbolicTensorSnapshot: shape: Tuple[Optional[str], ...] +@dataclass(frozen=True) +class _TensorMetadataSnapshot: + shape: Tuple[str, ...] + dtype: str + layout: str + device: str + requires_grad: bool + stride: Optional[Tuple[str, ...]] + + def _symbolic_scalar_snapshot( value: Argument, ) -> Optional[Tuple[str, str]]: @@ -146,6 +156,10 @@ def _extract_symbolic_snapshot(value: Argument) -> Any: return None +def _target_matches_by_identity(target: Any, targets: Tuple[Any, ...]) -> bool: + return any(target is candidate for candidate in targets) + + class NodeMetadata: def __init__(self, data: Dict[str, Any]) -> None: self.data: Dict[str, Any] = data.copy() @@ -231,6 +245,56 @@ class ExportPassBaseError(RuntimeError): pass +class _FastCopyFallback(Exception): + pass + + +def _unwrap_proxy_values(value: Argument) -> Argument: + while any( + isinstance(leaf, ProxyValue) for leaf in pytree.tree_leaves(value) + ): + value = pytree.tree_map_only(ProxyValue, lambda proxy: proxy.data, value) + return value + + +def _tensor_metadata_snapshot(value: torch.Tensor) -> _TensorMetadataSnapshot: + stride = None + if value.layout == torch.strided: + stride = tuple(str(dim) for dim in value.stride()) + return _TensorMetadataSnapshot( + shape=tuple(str(dim) for dim in value.shape), + dtype=str(value.dtype), + layout=str(value.layout), + device=str(value.device), + requires_grad=value.requires_grad, + stride=stride, + ) + + +def _tensor_metadata_changed(original: Argument, new: Argument) -> bool: + original = _unwrap_proxy_values(original) + new = _unwrap_proxy_values(new) + 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 _tensor_metadata_snapshot(original_leaf) != _tensor_metadata_snapshot( + new_leaf + ): + return True + + return False + + @dataclass(frozen=True) class ExportedProgramPassResult: exported_program: ExportedProgram @@ -283,6 +347,20 @@ def ensures(self, exported_program: ExportedProgram) -> None: # noqa: B027 """ +_FAST_COPY_UNSAFE_TARGETS: Tuple[Any, ...] = ( + torch.ops.aten.convolution, + torch.ops.aten.convolution.default, + torch.ops.aten.linear, + torch.ops.aten.linear.default, + exir_ops.edge.aten.convolution.default, + exir_ops.edge.aten.linear.default, +) + + +def _is_fast_copy_unsafe_target(target: Any) -> bool: + return _target_matches_by_identity(target, _FAST_COPY_UNSAFE_TARGETS) + + class _ExportPassBase(PassBase): """ Interpreter-based pass class to help users maintain the IR spec while writing @@ -416,12 +494,54 @@ def make_tensor_meta(x: Argument) -> Optional[TensorMetadata]: 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 one of the exact targets) 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_tuple = tuple(targeted) + except TypeError: + targeted_tuple = None + self._targeted_ops: Optional[Tuple[Any, ...]] = targeted_tuple + else: + self._targeted_ops: Optional[Tuple[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 + self._fast_copy_enabled: bool = ( + self._targeted_ops is not None and not call_overridden + ) + + # 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, @@ -515,10 +635,181 @@ def call_method( # pyre-fixme[14] ) -> None: raise ExportPassBaseError("call_method is not supported.") + # -- fast-copy helpers ------------------------------------------ + + def _fetch_attr_for_fast_copy(self, target: str) -> Any: + attr_itr = self.module + for atom in target.split("."): + try: + attr_itr = getattr(attr_itr, atom) + except AttributeError as exc: + raise _FastCopyFallback from exc + return attr_itr + + @staticmethod + def _preflight_get_attr_destinations( + tracer: "_ExportPassBase.ExportTracer", + get_attr_values: Dict[torch.fx.Node, Tuple[Any, List[str]]], + ) -> None: + planned_paths = { + tuple(target_atoms) for _, target_atoms in get_attr_values.values() + } + for path in planned_paths: + if len(path) == 0: + raise _FastCopyFallback + for index in range(1, len(path)): + if path[:index] in planned_paths: + raise _FastCopyFallback + + for value, target_atoms in get_attr_values.values(): + root = tracer.root + for atom in target_atoms[:-1]: + if not hasattr(root, atom): + break + child = getattr(root, atom) + if not isinstance(child, torch.nn.Module): + raise _FastCopyFallback + root = child + else: + leaf_name = target_atoms[-1] + if hasattr(root, leaf_name) and getattr(root, leaf_name) is not value: + raise _FastCopyFallback + + 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 = old_node.target + assert isinstance(target, str) + get_attr_values[old_node] = ( + self._fetch_attr_for_fast_copy(target), + target.split("."), + ) + self._preflight_get_attr_destinations(tracer, get_attr_values) + 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 + + 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 + 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 self._targeted_ops is not None + and not _target_matches_by_identity(n.target, self._targeted_ops) + and self.callback.should_fast_copy_node(n.target) + and n.meta.get("val") is not None + and "tensor_meta" in n.meta + ): + 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 _target_matches_by_identity(n.target, self._targeted_ops) + ): + if _tensor_metadata_changed(n.meta.get("val"), result): + self._fast_copy_enabled = False + + return result def __init__(self) -> None: self.interpreter = torch.fx.Interpreter( @@ -537,6 +828,14 @@ def should_preserve_symbolic_input_metadata(self) -> bool: """ return True + def should_fast_copy_node(self, target: torch.fx.node.Target) -> bool: + """Return whether a cold call_function node can bypass replay. + + Passes with subclass-wide ``call_operator`` behavior can override this + to keep selected non-targeted operators on the normal replay path. + """ + return not _is_fast_copy_unsafe_target(target) + def _capture_symbolic_input_snapshots( self, graph_module: fx.GraphModule ) -> List[Any]: @@ -823,13 +1122,17 @@ def output(self, results: List[Argument], meta: NodeMetadata) -> ProxyValue: 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(): @@ -879,6 +1182,14 @@ def call(self, graph_module: fx.GraphModule) -> PassResult: 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): diff --git a/exir/tests/test_pass_infra.py b/exir/tests/test_pass_infra.py index 16ed5af4180..1e164908d2f 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -8,6 +8,7 @@ # pyre-strict import unittest +from typing import Any import executorch.exir as exir import torch @@ -28,6 +29,7 @@ from torch.export import Dim, export, ExportedProgram from torch.export.graph_signature import InputKind, InputSpec, TensorArgument from torch.fx.passes.infra.pass_base import PassBase, PassResult +from torch.fx.passes.shape_prop import _extract_tensor_metadata class TestPassInfra(unittest.TestCase): @@ -229,6 +231,318 @@ def test_rejects_implicit_symbolic_scalar_coercions(self) -> None: float(ProxyValue(sym_float, torch.fx.Graph().placeholder("x"))) +class TestExportPassFastCopy(unittest.TestCase): + @staticmethod + def _edge_graph_module(module: torch.nn.Module) -> torch.fx.GraphModule: + return ( + to_edge(export(module, (torch.randn(2),), strict=True)) + .exported_program() + .graph_module + ) + + @staticmethod + def _raw_add_graph_module( + dynamic_shapes: Any | None = None, + ) -> torch.fx.GraphModule: + class AddModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + x + + return export( + AddModule(), + (torch.randn(2),), + dynamic_shapes=dynamic_shapes, + strict=True, + ).graph_module + + @staticmethod + def _ensure_tensor_meta(graph_module: torch.fx.GraphModule) -> None: + for node in graph_module.graph.nodes: + value = node.meta.get("val") + if isinstance(value, torch.Tensor) and "tensor_meta" not in node.meta: + node.meta["tensor_meta"] = _extract_tensor_metadata(value) + + @staticmethod + def _call_function_targets( + graph_module: torch.fx.GraphModule, + ) -> list[torch.fx.node.Target]: + return [ + node.target for node in graph_module.graph.nodes if node.op == "call_function" + ] + + def test_empty_targeted_ops_does_not_fall_back_to_target_ops(self) -> None: + class AddModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + x + + class EmptyTargetedOpsPass(ExportPass): + targeted_ops: tuple[()] = () + target_ops = {exir_ops.edge.aten.add.Tensor} + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_calls += 1 + return super().call_operator(op, args, kwargs, meta) + + graph_module = self._edge_graph_module(AddModule()) + pass_ = EmptyTargetedOpsPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 0) + + def test_missing_tensor_meta_uses_normal_replay(self) -> None: + graph_module = self._edge_graph_module(self._AddModule()) + add_node = self._single_call_function_node( + graph_module, exir_ops.edge.aten.add.Tensor + ) + del add_node.meta["tensor_meta"] + + class TargetedPass(ExportPass): + targeted_ops = (exir_ops.edge.aten.mul.Tensor,) + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_calls += 1 + return super().call_operator(op, args, kwargs, meta) + + pass_ = TargetedPass() + new_graph_module = pass_(graph_module).graph_module + new_add_node = self._single_call_function_node( + new_graph_module, exir_ops.edge.aten.add.Tensor + ) + + self.assertEqual(pass_.operator_calls, 1) + self.assertIn("tensor_meta", new_add_node.meta) + + def test_should_fast_copy_node_hook_keeps_selected_cold_ops_on_slow_path( + self, + ) -> None: + graph_module = self._edge_graph_module(self._AddModule()) + + class HookedPass(ExportPass): + targeted_ops: tuple[()] = () + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def should_fast_copy_node(self, target: torch.fx.node.Target) -> bool: + return target is not exir_ops.edge.aten.add.Tensor + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_calls += 1 + return super().call_operator(op, args, kwargs, meta) + + pass_ = HookedPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 1) + + def test_packet_target_does_not_match_overload_target(self) -> None: + graph_module = self._raw_add_graph_module() + self._ensure_tensor_meta(graph_module) + + class PacketTargetPass(ExportPass): + targeted_ops = (torch.ops.aten.add,) + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_calls += 1 + return super().call_operator(op, args, kwargs, meta) + + pass_ = PacketTargetPass() + new_graph_module = pass_(graph_module).graph_module + + self.assertEqual(pass_.operator_calls, 0) + self.assertEqual( + self._call_function_targets(new_graph_module), [torch.ops.aten.add.Tensor] + ) + + def test_symbolic_metadata_drift_check_does_not_force_symint_bool( + self, + ) -> None: + graph_module = self._raw_add_graph_module( + dynamic_shapes=({0: Dim("batch", min=1, max=8)},) + ) + + class SymbolicTargetPass(ExportPass): + targeted_ops = (torch.ops.aten.add.Tensor,) + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_calls += 1 + return super().call_operator(op, args, kwargs, meta) + + pass_ = SymbolicTargetPass() + new_graph_module = pass_(graph_module).graph_module + + self.assertEqual(pass_.operator_calls, 1) + self.assertEqual( + self._call_function_targets(new_graph_module), [torch.ops.aten.add.Tensor] + ) + + def test_nested_target_output_metadata_drift_disables_downstream_fast_copy( + self, + ) -> None: + class MaxThenAddModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + values, _ = torch.max(x, dim=1) + return values + values + + graph_module = export( + MaxThenAddModule(), + (torch.randn(2, 3),), + strict=True, + ).graph_module + self._ensure_tensor_meta(graph_module) + + class TupleMetadataDriftPass(ExportPass): + targeted_ops = (torch.ops.aten.max.dim,) + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue | tuple[ProxyValue, ProxyValue]: + self.operator_calls += 1 + result = super().call_operator(op, args, kwargs, meta) + if op is not torch.ops.aten.max.dim: + return result + + values = self.call_getitem(result, 0, meta) + indices = self.call_getitem(result, 1, meta) + return (ProxyValue(values.data.unsqueeze(0), values.proxy), indices) + + pass_ = TupleMetadataDriftPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 2) + + def test_overlapping_get_attr_fast_copy_fallback_is_atomic(self) -> None: + class TargetedPass(ExportPass): + targeted_ops = (torch.ops.aten.mul.Tensor,) + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_calls += 1 + return super().call_operator(op, args, kwargs, meta) + + weight = torch.ones(2) + weight.x = torch.ones(2) + root = torch.nn.Module() + root.register_buffer("w", weight) + graph = torch.fx.Graph() + w = graph.get_attr("w") + wx = graph.get_attr("w.x") + cold_node = graph.call_function(torch.ops.aten.add.Tensor, (w, wx)) + cold_node.meta["val"] = weight + weight.x + cold_node.meta["tensor_meta"] = _extract_tensor_metadata(cold_node.meta["val"]) + graph.output(cold_node) + graph_module = torch.fx.GraphModule(root, graph) + pass_ = TargetedPass() + + new_graph_module = pass_(graph_module).graph_module + + self.assertEqual(pass_.operator_calls, 1) + self.assertFalse( + any( + node.op == "get_attr" and len(node.users) == 0 + for node in new_graph_module.graph.nodes + ) + ) + + def test_unrelated_runtime_error_during_fast_copy_propagates(self) -> None: + graph_module = self._edge_graph_module(self._AddModule()) + + class RaisingFastCopyPass(ExportPass): + targeted_ops: tuple[()] = () + + class ExportInterpreter(ExportPass.ExportInterpreter): + def _fast_copy_node(self, n: torch.fx.Node) -> ProxyValue: + raise RuntimeError("unrelated fast-copy failure") + + with self.assertRaisesRegex(RuntimeError, "unrelated fast-copy failure"): + RaisingFastCopyPass()(graph_module) + + class _AddModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + x + + @staticmethod + def _single_call_function_node( + graph_module: torch.fx.GraphModule, + target: torch.fx.node.Target, + ) -> torch.fx.Node: + matches = [ + node + for node in graph_module.graph.nodes + if node.op == "call_function" and node.target is target + ] + if len(matches) != 1: + raise AssertionError(f"Expected exactly one {target} node, found {matches}") + return matches[0] + + class TestExportedProgramPassManager(unittest.TestCase): def test_runs_graph_module_passes_on_exported_program(self) -> None: """