From c10aed399b4033b3223a954e0d7fe9a409e86f2f Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Sun, 9 Aug 2026 17:28:22 -0700 Subject: [PATCH 1/2] Fast-copy cold nodes during ExportPass retracing Summary: After ARM pass-skipping and targeted-op ownership landed separately in D106781989, this diff optimizes the remaining `ExportPass` replay cost for passes that declare `target_ops` or `targeted_ops`. Cold operator nodes are copied with `graph.node_copy` instead of being re-dispatched through FakeTensor. Old-to-new node remapping preserves dependencies and `get_attr` values. Fast-copy is disabled when `call()` is overridden, for exact convolution or linear targets, or after a hot node changes nested tensor metadata. The fast path preflights every input before mutating the new graph or module tree, so a remapping fallback cannot leave orphaned nodes, attributes, or remap entries. An explicitly empty `targeted_ops` remains authoritative instead of falling back to legacy `target_ops`. Nested ARM control-flow submodules use the established `ArmPass.should_run_pass()` contract. This diff does not duplicate pass auto-discovery or generic skipping logic owned by D106781989. Per review from `jrstevens`, the process-global FakeTensor cache extension is now isolated in child D115374097 so its monkeypatching design and incremental performance can be reviewed independently. A/B benchmark against the parent revision on the same host, each side run twice: - CombinedControl U55 lowering: 148.465s / 149.726s with fast-copy vs 165.886s / 171.181s before it. Warm speedup: 12.5%; two-run mean speedup: 11.5%. - Synthetic U55 suite: 81.238s / 82.371s vs 83.816s / 80.802s. The difference is within run-to-run noise; the large model is the representative workload. Differential Revision: D97528110 --- backends/arm/_passes/arm_pass.py | 2 + exir/pass_base.py | 266 ++++++++++++++++++++++++++++++- exir/tests/test_pass_infra.py | 87 ++++++++++ 3 files changed, 350 insertions(+), 5 deletions(-) diff --git a/backends/arm/_passes/arm_pass.py b/backends/arm/_passes/arm_pass.py index 8afb8cc6e1f..9bc6932b7e5 100644 --- a/backends/arm/_passes/arm_pass.py +++ b/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. diff --git a/exir/pass_base.py b/exir/pass_base.py index c657ac53a91..bac16af40bf 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -143,6 +143,37 @@ def _extract_symbolic_snapshot(value: Argument) -> Any: 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 @@ class ExportPassBaseError(RuntimeError): pass +class _FastCopyFallback(RuntimeError): + pass + + @dataclass(frozen=True) class ExportedProgramPassResult: exported_program: ExportedProgram @@ -280,6 +315,22 @@ def ensures(self, exported_program: ExportedProgram) -> None: # noqa: B027 """ +_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 @@ 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* 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 @@ def call_method( # pyre-fixme[14] ) -> 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 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(): @@ -824,6 +1072,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 59406b13f8f..f8802483200 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 @@ -228,6 +229,92 @@ def test_rejects_implicit_symbolic_scalar_coercions(self) -> None: float(ProxyValue(sym_float, torch.fx.Graph().placeholder("x"))) +class TestExportPassFastCopy(unittest.TestCase): + 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: set[object] = set() + 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 = ( + to_edge(export(AddModule(), (torch.randn(2),), strict=True)) + .exported_program() + .graph_module + ) + pass_ = EmptyTargetedOpsPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 0) + + def test_fast_copy_fallback_is_side_effect_free(self) -> None: + class RootModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.branch = torch.nn.Module() + self.branch.register_buffer("weight", torch.ones(2)) + + root = RootModule() + graph = torch.fx.Graph() + x = graph.placeholder("x") + weight = graph.get_attr("branch.weight") + unresolved = graph.call_function(torch.ops.aten.neg.default, (x,)) + cold_node = graph.call_function(torch.ops.aten.add.Tensor, (weight, unresolved)) + graph.output(cold_node) + graph_module = torch.fx.GraphModule(root, graph) + + class TargetedPass(ExportPass): + targeted_ops = {torch.ops.aten.mul.Tensor} + + pass_ = TargetedPass() + pass_.tracer = pass_.ExportTracer(pass_, graph_module.graph._codegen) + interpreter = pass_.ExportInterpreter(pass_, graph_module) + + with self.assertRaises(RuntimeError): + interpreter._fast_copy_node(cold_node) + + self.assertEqual(list(pass_.tracer.graph.nodes), []) + self.assertFalse(hasattr(pass_.tracer.root, "branch")) + self.assertEqual(interpreter._node_remap, {}) + + def test_fast_copy_falls_back_for_unmapped_placeholder(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + cold_node = graph.call_function(torch.ops.aten.add.Tensor, (x, x)) + graph.output(cold_node) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + class TargetedPass(ExportPass): + targeted_ops = {torch.ops.aten.mul.Tensor} + + pass_ = TargetedPass() + pass_.tracer = pass_.ExportTracer(pass_, graph_module.graph._codegen) + interpreter = pass_.ExportInterpreter(pass_, graph_module) + + with self.assertRaises(RuntimeError): + interpreter._fast_copy_node(cold_node) + + self.assertEqual(list(pass_.tracer.graph.nodes), []) + self.assertEqual(interpreter._node_remap, {}) + + class TestExportedProgramPassManager(unittest.TestCase): def test_runs_graph_module_passes_on_exported_program(self) -> None: """ From 6caf0611b60ebf5233ace7aae0525ec14a3d1745 Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Sun, 9 Aug 2026 17:28:22 -0700 Subject: [PATCH 2/2] Add targeted_ops to core ExportPass passes for fast-copy optimization (#21702) Summary: Annotate four core `ExportPass` subclasses across three files with `targeted_ops` so they participate in the cold-node fast-copy path introduced by parent D97528110. When fast-copy remains safe, D975 copies eligible non-targeted `call_function` nodes with existing `val` metadata instead of replaying them through FakeTensor dispatch. This diff does not add generic pass skipping: every pass still enters replay, and hot nodes continue through each existing `call_operator` implementation. Each set exactly mirrors the rewrite guard already present in its pass: - `MemoryFormatOpsPass` targets `DimOrderOpsMap` keys. - `DimOrderOpsRevertPass` targets `MemoryFormatOpsMap` keys. - `NormalizeTransposePass` targets `aten.t.default`. - `RemoveMixedTypeOperators` targets `add.Tensor`, `mul.Tensor`, `sub.Tensor`, both tensor `div` overloads, and `minimum.default`. The diff is now stacked directly on D97528110, which owns the consumer infrastructure. Its previous D109215317 parent left these attributes inert. Direct raw-graph regressions prove that transpose normalization and mixed-type rewriting still dispatch hot operators while eligible unrelated operators bypass `call_operator`; the mixed-type test also verifies inserted promotion casts and output dtypes. A/B against D97528110 on the same host, with each revision run twice: - Synthetic add/conv/fc/mul U55 suite: 79.000s / 79.516s here versus 81.238s / 82.371s in D975. Mean speedup 3.1%; warmed speedup 3.5%. - CombinedControl U55 LOWERING: 145.801s / 145.888s here versus 148.465s / 149.726s in D975. Mean speedup 2.2%; warmed speedup 2.6%. Differential Revision: D96247716 --- exir/passes/memory_format_ops_pass.py | 10 ++- exir/passes/normalize_transpose_pass.py | 5 +- exir/passes/remove_mixed_type_operators.py | 9 +++ exir/tests/test_pass_infra.py | 93 ++++++++++++++++++++++ 4 files changed, 113 insertions(+), 4 deletions(-) diff --git a/exir/passes/memory_format_ops_pass.py b/exir/passes/memory_format_ops_pass.py index 13468dfd8d8..8bd260f1d57 100644 --- a/exir/passes/memory_format_ops_pass.py +++ b/exir/passes/memory_format_ops_pass.py @@ -29,6 +29,8 @@ class MemoryFormatOpsPass(ExportPass): 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 @@ def call_operator(self, op, args, kwargs, meta): 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 @@ class DimOrderOpsRevertPass(ExportPass): 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( diff --git a/exir/passes/normalize_transpose_pass.py b/exir/passes/normalize_transpose_pass.py index b401d2ad9b7..f062accd31a 100644 --- a/exir/passes/normalize_transpose_pass.py +++ b/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( diff --git a/exir/passes/remove_mixed_type_operators.py b/exir/passes/remove_mixed_type_operators.py index 86a71354337..058993ab1d1 100644 --- a/exir/passes/remove_mixed_type_operators.py +++ b/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: diff --git a/exir/tests/test_pass_infra.py b/exir/tests/test_pass_infra.py index f8802483200..f78f240ff8c 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -23,7 +23,11 @@ ) from executorch.exir.pass_manager import ExportedProgramPassManager, PassManager from executorch.exir.passes import ScalarToTensorPass +from executorch.exir.passes.normalize_transpose_pass import NormalizeTransposePass from executorch.exir.passes.pass_registry import PassRegistry +from executorch.exir.passes.remove_mixed_type_operators import ( + RemoveMixedTypeOperators, +) from executorch.exir.program import to_edge from torch.export import Dim, export, ExportedProgram from torch.export.graph_signature import InputKind, InputSpec, TensorArgument @@ -264,6 +268,95 @@ def call_operator( self.assertEqual(pass_.operator_calls, 0) + def test_normalize_transpose_uses_targeted_fast_copy(self) -> None: + class TransposeModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.t(torch.neg(x)) + + class RecordingNormalizeTransposePass(NormalizeTransposePass): + def __init__(self) -> None: + super().__init__() + self.operator_targets: list[Any] = [] + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_targets.append(op) + return super().call_operator(op, args, kwargs, meta) + + graph_module = export( + TransposeModule(), (torch.randn(2, 3),), strict=True + ).graph_module + pass_ = RecordingNormalizeTransposePass() + + pass_result = pass_(graph_module) + assert pass_result is not None + targets = [ + node.target + for node in pass_result.graph_module.graph.nodes + if node.op == "call_function" + ] + + self.assertEqual(pass_.operator_targets, [torch.ops.aten.t.default]) + self.assertIn(torch.ops.aten.t_copy.default, targets) + self.assertIn(torch.ops.aten.neg.default, targets) + + def test_remove_mixed_type_operators_uses_targeted_fast_copy(self) -> None: + class MixedTypeModule(torch.nn.Module): + def forward( + self, x: torch.Tensor, y: torch.Tensor + ) -> torch.Tensor: + return torch.neg(torch.add(x, y)) + + class RecordingRemoveMixedTypeOperators(RemoveMixedTypeOperators): + def __init__(self) -> None: + super().__init__() + self.operator_targets: list[Any] = [] + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_targets.append(op) + return super().call_operator(op, args, kwargs, meta) + + graph_module = export( + MixedTypeModule(), + ( + torch.tensor([1, 2], dtype=torch.int64), + torch.tensor([1.0, 2.0], dtype=torch.float32), + ), + strict=True, + ).graph_module + pass_ = RecordingRemoveMixedTypeOperators() + + pass_result = pass_(graph_module) + assert pass_result is not None + graph = pass_result.graph_module.graph + targets = [ + node.target for node in graph.nodes if node.op == "call_function" + ] + + self.assertIn(torch.ops.aten.add.Tensor, pass_.operator_targets) + self.assertIn(torch.ops.aten._to_copy.default, pass_.operator_targets) + self.assertNotIn(torch.ops.aten.neg.default, pass_.operator_targets) + self.assertIn(torch.ops.aten.neg.default, targets) + + add_nodes = graph.find_nodes( + op="call_function", target=torch.ops.aten.add.Tensor + ) + self.assertEqual(len(add_nodes), 1) + for arg in add_nodes[0].args: + assert isinstance(arg, torch.fx.Node) + self.assertEqual(arg.meta["val"].dtype, torch.float32) + def test_fast_copy_fallback_is_side_effect_free(self) -> None: class RootModule(torch.nn.Module): def __init__(self) -> None: