diff --git a/backends/arm/_passes/arm_pass.py b/backends/arm/_passes/arm_pass.py index 5c3541a7586..38a8796b92b 100644 --- a/backends/arm/_passes/arm_pass.py +++ b/backends/arm/_passes/arm_pass.py @@ -24,6 +24,12 @@ from torch.utils import _pytree as pytree +_OPS_WITHOUT_QUANTIZED_FAKE_KERNEL = ( + exir_ops.edge.aten.bmm.default, + exir_ops.edge.aten.leaky_relu.default, +) + + class ArmPass(ExportPass): """Base class for Arm passes.""" @@ -57,6 +63,13 @@ 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: + # Metadata is unavailable here, so conservatively replay every op whose + # quantized path needs a custom fake result. + 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]]: @@ -86,12 +99,8 @@ def get_name(pass_) -> str: ) def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False): - ops_without_quantized_fake_kernel = { - exir_ops.edge.aten.bmm.default, - exir_ops.edge.aten.leaky_relu.default, - } if ( - op in ops_without_quantized_fake_kernel + op in _OPS_WITHOUT_QUANTIZED_FAKE_KERNEL and isinstance(meta, NodeMetadata) and len(meta.data.get("input_qparams", {})) > 0 ): @@ -138,21 +147,32 @@ def _call_quantized_op_without_fake_kernel( self.tracer.set_metadata(res_proxy.node, res_data) return ProxyValue(res_data, res_proxy) + def should_run_on_nested_submodule(self, graph_module: GraphModule) -> bool: + """Return whether subclass rewrites should run in this nested submodule.""" + has_custom_precheck = type(self).should_run_pass is not ArmPass.should_run_pass + return has_custom_precheck and self.should_run_pass(graph_module) + def call_submodule( self, graph_module: GraphModule, inputs: tuple[Any, ...] ) -> PassResult: self.submodule_depth += 1 - if self.submodule_depth == 1: - 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. - _call_operator_fn = self.call_operator - self.call_operator = super().call_operator # type: ignore - result = super().call_submodule(graph_module, inputs) - self.call_operator = _call_operator_fn # type: ignore - self.submodule_depth -= 1 - return result + try: + if self.submodule_depth == 1: + self._top_level_call_operator = self.call_operator + return super().call_submodule(graph_module, inputs) + + call_operator = self.call_operator + try: + if self.should_run_on_nested_submodule(graph_module): + self.call_operator = self._top_level_call_operator # type: ignore + else: + # Nested submodules still need replay without subclass rewrites. + self.call_operator = super().call_operator # type: ignore + return super().call_submodule(graph_module, inputs) + finally: + self.call_operator = call_operator # type: ignore + finally: + self.submodule_depth -= 1 def call_shape_operator( self, op, args: tuple, kwargs: dict, meta: NodeMetadata, updated: bool = True @@ -235,10 +255,15 @@ def __call__(self, graph_module: GraphModule) -> PassResult | None: class ArmOpTargetedPass(ArmPass): """Base class for passes that only transform selected operators. - Subclasses set ``target_ops`` to the call_function targets they can - transform. If the current graph and nested control-flow subgraphs do not - contain any target, the pass returns immediately without paying the default - ExportPass retracing cost. + Subclasses must set ``target_ops`` to the exhaustive set of call_function + targets their ``call_operator()`` can transform. This ARM-specific contract + drives both the target pre-scan and fast-copy eligibility; an empty + ``target_ops`` disables fast copy and skips the pass. Generic ``ExportPass`` + subclasses instead use ``targeted_ops`` for explicit fast-copy opt-in. + + If the current graph and nested control-flow subgraphs do not contain any + target, the pass returns immediately without paying the default ExportPass + retracing cost. Set ``check_allowed_to_transform`` to ``True`` when the target pre-scan should also apply ``allowed_to_transform()`` to matching target nodes. This @@ -248,14 +273,28 @@ class ArmOpTargetedPass(ArmPass): """ + enable_fast_copy = True target_ops: Collection[Any] = () check_allowed_to_transform = False - def has_target_node(self, graph_module: GraphModule) -> bool: + def get_fast_copy_target_ops(self) -> Optional[tuple[Any, ...]]: + if not self.enable_fast_copy: + return None + targets = tuple(self.target_ops) + return targets if targets else None + + def should_run_on_nested_submodule(self, graph_module: GraphModule) -> bool: + return self.has_target_node(graph_module, recursive=False) + + def has_target_node( + self, graph_module: GraphModule, recursive: bool = True + ) -> bool: """Return whether the graph module tree contains a target node. Args: graph_module (GraphModule): The graph module tree to inspect. + recursive (bool): Whether to inspect nested child GraphModules. + When false, inspect only ``graph_module`` itself. Returns: bool: True if a matching call_function node is present. @@ -284,7 +323,7 @@ def graph_has_target(module: GraphModule) -> bool: if target_node_can_trigger_pass(node): return True - return any( + return recursive and any( isinstance(child, GraphModule) and graph_has_target(child) for child in module.children() ) diff --git a/backends/arm/test/passes/test_arm_op_targeted_pass.py b/backends/arm/test/passes/test_arm_op_targeted_pass.py index e990e13bb08..3c3bf1cfd54 100644 --- a/backends/arm/test/passes/test_arm_op_targeted_pass.py +++ b/backends/arm/test/passes/test_arm_op_targeted_pass.py @@ -15,6 +15,7 @@ from executorch.exir.pass_base import ExportPass from torch.fx import Graph, GraphModule from torch.fx.passes.infra.pass_base import PassResult +from torch.fx.passes.shape_prop import _extract_tensor_metadata TARGET_OP = torch.ops.aten.add.Tensor @@ -50,15 +51,17 @@ def run_single_pass(graph_module: GraphModule, test_pass: ExportPass) -> PassRes class DummyTargetedPass(ArmOpTargetedPass): _passes_required_after: Set[Type[ExportPass]] = set() - target_ops = (TARGET_OP,) + target_ops: tuple[torch._ops.OpOverload, ...] = (TARGET_OP,) check_allowed_to_transform = True def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.call_operator_count = 0 + self.called_ops = [] def call_operator(self, op, args, kwargs, meta): self.call_operator_count += 1 + self.called_ops.append(op) return super().call_operator(op, args, kwargs, meta) @@ -80,6 +83,28 @@ def call(self, graph_module: GraphModule) -> PassResult: return PassResult(graph_module, True) +class NestedCondModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + def outer_true(arg: torch.Tensor) -> torch.Tensor: + def inner_true(inner_arg: torch.Tensor) -> torch.Tensor: + return inner_arg + inner_arg + + def inner_false(inner_arg: torch.Tensor) -> torch.Tensor: + return inner_arg * inner_arg + + return torch.cond( + arg.sum() > 1, + inner_true, + inner_false, + [arg], + ) + + def outer_false(arg: torch.Tensor) -> torch.Tensor: + return arg * arg + + return torch.cond(x.sum() > 0, outer_true, outer_false, [x]) + + class CondModule(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: def true_branch(arg: torch.Tensor) -> torch.Tensor: @@ -91,6 +116,43 @@ def false_branch(arg: torch.Tensor) -> torch.Tensor: return torch.cond(x.sum() > 0, true_branch, false_branch, [x]) +def test_replays_targeted_op_and_fast_copies_cold_op() -> None: + graph = Graph() + lhs = graph.placeholder("lhs") + rhs = graph.placeholder("rhs") + lhs.meta["val"] = torch.randn(2, 3) + rhs.meta["val"] = torch.randn(2, 3) + cold = graph.call_function(torch.ops.aten.mul.Tensor, (lhs, rhs)) + cold.meta["val"] = lhs.meta["val"] * rhs.meta["val"] + cold.meta["tensor_meta"] = _extract_tensor_metadata(cold.meta["val"]) + targeted = graph.call_function(TARGET_OP, (cold, rhs)) + targeted.meta["val"] = cold.meta["val"] + rhs.meta["val"] + targeted.meta["tensor_meta"] = _extract_tensor_metadata(targeted.meta["val"]) + graph.output(targeted) + graph_module = GraphModule(torch.nn.Module(), graph) + targeted_pass = DummyTargetedPass() + + result = run_single_pass(graph_module, targeted_pass) + + assert result.modified + assert targeted_pass.call_operator_count == 1 + + +def test_empty_target_ops_disables_fast_copy_and_skips_pass() -> None: + class EmptyTargetedPass(DummyTargetedPass): + target_ops: tuple[torch._ops.OpOverload, ...] = () + + graph_module = create_graph_module(TARGET_OP) + targeted_pass = EmptyTargetedPass() + + result = run_single_pass(graph_module, targeted_pass) + + assert targeted_pass.get_fast_copy_target_ops() is None + assert result.graph_module is graph_module + assert not result.modified + assert targeted_pass.call_operator_count == 0 + + def test_skips_when_target_is_absent() -> None: graph_module = create_graph_module() targeted_pass = DummyTargetedPass() @@ -138,6 +200,18 @@ def test_runs_when_previous_pass_creates_target() -> None: assert targeted_pass.call_operator_count == 1 +def test_runs_when_target_is_present_in_deeply_nested_submodule() -> None: + exported_program = torch.export.export(NestedCondModule(), (torch.randn(2, 3),)) + graph_module = exported_program.graph_module + targeted_pass = DummyTargetedPass() + + result = run_single_pass(graph_module, targeted_pass) + + assert result is not None + assert result.modified + assert TARGET_OP in targeted_pass.called_ops + + def test_runs_when_target_is_present_in_nested_submodule() -> None: exported_program = torch.export.export(CondModule(), (torch.randn(2, 3),)) graph_module = exported_program.graph_module @@ -147,4 +221,4 @@ def test_runs_when_target_is_present_in_nested_submodule() -> None: assert result is not None assert result.modified - assert targeted_pass.call_operator_count > 0 + assert TARGET_OP in targeted_pass.called_ops diff --git a/exir/pass_base.py b/exir/pass_base.py index 6071aae2be8..e4c72ffcd15 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -98,11 +98,21 @@ def _unstack_pytree(xs) -> List[PyTree]: # pyre-ignore return pytrees -@dataclass(frozen=True) +@dataclass(frozen=True, slots=True) class _SymbolicTensorSnapshot: shape: Tuple[Optional[str], ...] +@dataclass(frozen=True, slots=True) +class _TensorMetadataSnapshot: + shape: Tuple[Any, ...] + dtype: torch.dtype + layout: torch.layout + device: torch.device + requires_grad: bool + stride: Optional[Tuple[Any, ...]] + + 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,86 @@ class ExportPassBaseError(RuntimeError): pass +class _FastCopyFallback(Exception): + pass + + +def _metadata_dimension_snapshot(value: Any) -> Any: + symbolic = _symbolic_scalar_snapshot(value) + return symbolic if symbolic is not None else value + + +def _tensor_metadata_snapshot(value: torch.Tensor) -> _TensorMetadataSnapshot: + stride = None + if value.layout == torch.strided: + stride = tuple(_metadata_dimension_snapshot(dim) for dim in value.stride()) + return _TensorMetadataSnapshot( + shape=tuple(_metadata_dimension_snapshot(dim) for dim in value.shape), + dtype=value.dtype, + layout=value.layout, + device=value.device, + requires_grad=value.requires_grad, + stride=stride, + ) + + +def _metadata_leaf_snapshot(value: Any) -> Tuple[type, Any]: + symbolic = _symbolic_scalar_snapshot(value) + if symbolic is not None: + return (type(value), symbolic) + if isinstance( + value, + ( + type(None), + bool, + int, + float, + complex, + str, + bytes, + torch.dtype, + torch.layout, + torch.device, + torch.memory_format, + ), + ): + return (type(value), value) + # Avoid arbitrary equality/repr work. Distinct unknown objects may be + # equivalent, but treating them as drift is the conservative choice. + return (type(value), id(value)) + + +def _tensor_metadata_changed(original: Argument, new: Argument) -> bool: + if isinstance(original, ProxyValue): + original = original.data + if isinstance(new, ProxyValue): + new = new.data + 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): + if isinstance(original_leaf, ProxyValue): + original_leaf = original_leaf.data + if isinstance(new_leaf, ProxyValue): + new_leaf = new_leaf.data + + 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 original_is_tensor: + if _tensor_metadata_snapshot(original_leaf) != _tensor_metadata_snapshot( + new_leaf + ): + return True + elif _metadata_leaf_snapshot(original_leaf) != _metadata_leaf_snapshot(new_leaf): + return True + + return False + + @dataclass(frozen=True) class ExportedProgramPassResult: exported_program: ExportedProgram @@ -283,12 +377,53 @@ def ensures(self, exported_program: ExportedProgram) -> None: # noqa: B027 """ +# Replaying convolution and linear operators can refresh layout-sensitive +# metadata that downstream passes rely on. Keep every ATen convolution spelling +# that can survive export, plus the Edge aliases exposed here, on the replay path. +_FAST_COPY_UNSAFE_TARGETS: Tuple[Any, ...] = ( + torch.ops.aten.convolution, + torch.ops.aten.convolution.default, + torch.ops.aten.conv1d, + torch.ops.aten.conv1d.default, + torch.ops.aten.conv1d.padding, + torch.ops.aten.conv2d, + torch.ops.aten.conv2d.default, + torch.ops.aten.conv2d.padding, + torch.ops.aten.conv3d, + torch.ops.aten.conv3d.default, + torch.ops.aten.conv3d.padding, + torch.ops.aten.conv_transpose1d, + torch.ops.aten.conv_transpose1d.default, + torch.ops.aten.conv_transpose2d, + torch.ops.aten.conv_transpose2d.input, + torch.ops.aten.conv_transpose3d, + torch.ops.aten.conv_transpose3d.input, + torch.ops.aten.linear, + torch.ops.aten.linear.default, + exir_ops.edge.aten.convolution.default, + exir_ops.edge.aten.conv2d.default, + exir_ops.edge.aten.conv2d.padding, + exir_ops.edge.aten.conv3d.default, + exir_ops.edge.aten.conv3d.padding, + exir_ops.edge.aten.linear.default, +) +_FAST_COPY_UNSAFE_TARGET_IDS = frozenset( + id(target) for target in _FAST_COPY_UNSAFE_TARGETS +) + + +def _is_fast_copy_unsafe_target(target: Any) -> bool: + return id(target) in _FAST_COPY_UNSAFE_TARGET_IDS + + class _ExportPassBase(PassBase): """ Interpreter-based pass class to help users maintain the IR spec while writing transformations. """ + enable_fast_copy = False + @staticmethod def _create_dummy_node_metadata() -> NodeMetadata: return NodeMetadata({"stack_trace": "".join(traceback.format_stack(limit=1))}) @@ -300,7 +435,8 @@ def __init__(self, callback: "_ExportPassBase", codegen: CodeGen) -> None: self.root = torch.nn.Module() self.graph = torch.fx.Graph() self.graph.set_codegen(codegen) - self.tensor_attrs: Dict[str, torch.Tensor] = {} # type: ignore[assignment] + # PythonKeyTracer.create_arg expects tensor -> qualified name. + self.tensor_attrs: Dict[torch.Tensor, str] = {} self.fake_tensor_mode: Optional[FakeTensorMode] = None self.submodules: Dict[torch.nn.Module, str] = {} @@ -416,12 +552,44 @@ 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. + self._targeted_ops = callback.get_fast_copy_target_ops() + + # 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 +683,300 @@ def call_method( # pyre-fixme[14] ) -> None: raise ExportPassBaseError("call_method is not supported.") + # -- fast-copy helpers ------------------------------------------ + + @staticmethod + def _proxy_value_node( + value: Any, + tracer: "_ExportPassBase.ExportTracer", + ) -> Optional[torch.fx.Node]: + if not isinstance(value, ProxyValue): + return None + proxy_or_node = value.proxy_or_node + node = ( + proxy_or_node.node + if isinstance(proxy_or_node, torch.fx.Proxy) + else proxy_or_node + ) + if not isinstance(node, torch.fx.Node) or node.graph is not tracer.graph: + raise _FastCopyFallback + return node + + @staticmethod + def _source_attr_registration( + parent: torch.nn.Module, + name: str, + ) -> Tuple[str, bool]: + if name in parent._parameters: + return ("parameter", True) + if name in parent._buffers: + return ("buffer", name not in parent._non_persistent_buffers_set) + if name in parent._modules: + return ("module", True) + return ("attribute", True) + + def _fetch_attr_for_fast_copy( + self, target: str + ) -> Tuple[Any, Tuple[str, ...], str, bool]: + target_atoms = tuple(target.split(".")) + if not target_atoms or any(not atom for atom in target_atoms): + raise _FastCopyFallback + + parent = self.module + for atom in target_atoms[:-1]: + try: + parent = getattr(parent, atom) + except AttributeError as exc: + raise _FastCopyFallback from exc + if not isinstance(parent, torch.nn.Module): + raise _FastCopyFallback + + try: + value = getattr(parent, target_atoms[-1]) + except AttributeError as exc: + raise _FastCopyFallback from exc + registration, persistent = self._source_attr_registration( + parent, target_atoms[-1] + ) + return value, target_atoms, registration, persistent + + @staticmethod + def _destination_attr_matches( + parent: torch.nn.Module, + name: str, + value: Any, + registration: str, + persistent: bool, + ) -> bool: + if not hasattr(parent, name) or getattr(parent, name) is not value: + return False + if registration == "parameter": + return name in parent._parameters + if registration == "buffer": + return ( + name in parent._buffers + and (name not in parent._non_persistent_buffers_set) == persistent + ) + if registration == "module": + return name in parent._modules + return ( + name not in parent._parameters + and name not in parent._buffers + and name not in parent._modules + ) + + def _preflight_get_attr_destinations( + self, + tracer: "_ExportPassBase.ExportTracer", + get_attr_values: Dict[ + torch.fx.Node, Tuple[Any, Tuple[str, ...], str, bool] + ], + ) -> None: + planned: Dict[Tuple[str, ...], Tuple[Any, str, bool]] = {} + for value, target_atoms, registration, persistent in get_attr_values.values(): + previous = planned.get(target_atoms) + if previous is not None and ( + previous[0] is not value + or previous[1] != registration + or previous[2] != persistent + ): + raise _FastCopyFallback + planned[target_atoms] = (value, registration, persistent) + + for path in planned: + for index in range(1, len(path)): + if path[:index] in planned: + raise _FastCopyFallback + + for path, (value, registration, persistent) in planned.items(): + parent = tracer.root + for atom in path[:-1]: + if not hasattr(parent, atom): + break + child = getattr(parent, atom) + if not isinstance(child, torch.nn.Module): + raise _FastCopyFallback + parent = child + else: + leaf_name = path[-1] + if hasattr(parent, leaf_name) and not self._destination_attr_matches( + parent, + leaf_name, + value, + registration, + persistent, + ): + raise _FastCopyFallback + + def _preflight_fast_copy_inputs( + self, + n: torch.fx.Node, + tracer: "_ExportPassBase.ExportTracer", + ) -> Tuple[ + Dict[torch.fx.Node, torch.fx.Node], + Dict[torch.fx.Node, Tuple[Any, Tuple[str, ...], str, bool]], + ]: + resolved_nodes: Dict[torch.fx.Node, torch.fx.Node] = {} + get_attr_values: Dict[ + torch.fx.Node, Tuple[Any, Tuple[str, ...], str, bool] + ] = {} + for old_node in n.all_input_nodes: + new_node = self._node_remap.get(old_node) + if new_node is not None: + if new_node.graph is not tracer.graph: + raise _FastCopyFallback + resolved_nodes[old_node] = new_node + continue + + env_node = self._proxy_value_node(self.env.get(old_node), tracer) + if env_node is not None: + resolved_nodes[old_node] = env_node + continue + if old_node.op != "get_attr" or not isinstance(old_node.target, str): + raise _FastCopyFallback + + get_attr_values[old_node] = self._fetch_attr_for_fast_copy( + old_node.target + ) + + self._preflight_get_attr_destinations(tracer, get_attr_values) + return resolved_nodes, get_attr_values + + @staticmethod + def _get_attr_parent( + root: torch.nn.Module, + target_atoms: Tuple[str, ...], + ) -> torch.nn.Module: + parent = root + for atom in target_atoms[:-1]: + if not hasattr(parent, atom): + parent.add_module(atom, torch.nn.Module()) + child = getattr(parent, atom) + if not isinstance(child, torch.nn.Module): + raise ExportPassBaseError( + f"Cannot install get_attr through non-module attribute {atom}." + ) + parent = child + return parent + + @staticmethod + def _install_attr( + parent: torch.nn.Module, + name: str, + value: Any, + registration: str, + persistent: bool, + ) -> None: + if ( + name in parent.__dict__ + or name in parent._parameters + or name in parent._buffers + or name in parent._modules + ): + return + if registration == "parameter": + parent.register_parameter(name, value) + elif registration == "buffer": + parent.register_buffer(name, value, persistent=persistent) + elif registration == "module": + parent.add_module(name, value) + else: + setattr(parent, name, value) + + def _commit_fast_copy_get_attrs( + self, + tracer: "_ExportPassBase.ExportTracer", + resolved_nodes: Dict[torch.fx.Node, torch.fx.Node], + get_attr_values: Dict[ + torch.fx.Node, Tuple[Any, Tuple[str, ...], str, bool] + ], + ) -> None: + for old_node, ( + value, + target_atoms, + registration, + persistent, + ) in get_attr_values.items(): + parent = self._get_attr_parent(tracer.root, target_atoms) + self._install_attr( + parent, + target_atoms[-1], + value, + registration, + persistent, + ) + copied = tracer.graph.node_copy(old_node, lambda node: resolved_nodes[node]) + proxy_value = ProxyValue(value, torch.fx.Proxy(copied, tracer)) + if isinstance(value, torch.Tensor): + tracer.tensor_attrs[value] = str(old_node.target) + tracer.set_metadata(copied, value) + tracer.callback.on_attr(proxy_value) + resolved_nodes[old_node] = copied + self._node_remap[old_node] = copied + + def _fast_copy_node(self, n: torch.fx.Node) -> "ProxyValue": + tracer = self.callback.tracer + resolved_nodes, get_attr_values = self._preflight_fast_copy_inputs( + n, tracer + ) + self._commit_fast_copy_get_attrs(tracer, resolved_nodes, get_attr_values) + + new_node = tracer.graph.node_copy( + n, lambda old_node: resolved_nodes[old_node] + ) + val = n.meta.get("val") + result = ProxyValue(val, torch.fx.Proxy(new_node, tracer)) + self._node_remap[n] = new_node + return result + + def _record_slow_path_result(self, n: torch.fx.Node, result: Argument) -> None: + result_node = None + for leaf in pytree.tree_leaves(result): + try: + leaf_node = self._proxy_value_node(leaf, self.callback.tracer) + except _FastCopyFallback: + self._fast_copy_enabled = False + break + if leaf is result: + result_node = leaf_node + else: + if result_node is not None: + self._node_remap[n] = result_node + + if "val" in n.meta and _tensor_metadata_changed(n.meta["val"], result): + self._fast_copy_enabled = False + 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_copied = False + + # 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: + result = self._fast_copy_node(n) + fast_copied = True + except _FastCopyFallback: + result = super().run_node(n) + else: + result = super().run_node(n) + + if not fast_copied and self._fast_copy_enabled: + self._record_slow_path_result(n, result) + + return result def __init__(self) -> None: self.interpreter = torch.fx.Interpreter( @@ -537,6 +995,26 @@ def should_preserve_symbolic_input_metadata(self) -> bool: """ return True + def get_fast_copy_target_ops(self) -> Optional[Tuple[Any, ...]]: + """Return exact targets only when this pass explicitly enables fast-copy.""" + if not self.enable_fast_copy: + return None + targeted_ops = getattr(self, "targeted_ops", None) + if targeted_ops is None: + return None + try: + return tuple(targeted_ops) + except TypeError: + return None + + 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 +1301,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 +1361,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..f5887822a47 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,830 @@ def test_rejects_implicit_symbolic_scalar_coercions(self) -> None: float(ProxyValue(sym_float, torch.fx.Graph().placeholder("x"))) +class TestExportPassFastCopy(unittest.TestCase): + class _CountingTargetedPass(ExportPass): + enable_fast_copy = True + + def __init__( + self, + targeted_ops: tuple[Any, ...] = (torch.ops.aten.mul.Tensor,), + ) -> None: + super().__init__() + self.targeted_ops = targeted_ops + 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) + + @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_target_ops_alone_does_not_enable_fast_copy(self) -> None: + graph_module = self._edge_graph_module(self._AddModule()) + + class TargetOpsOnlyPass(ExportPass): + target_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_ = TargetOpsOnlyPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 1) + + def test_targeted_ops_alone_does_not_enable_fast_copy(self) -> None: + graph_module = self._edge_graph_module(self._AddModule()) + + class TargetedOpsOnlyPass(ExportPass): + targeted_ops: tuple[()] = () + + 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_ = TargetedOpsOnlyPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 1) + + def test_explicit_empty_targeted_ops_enables_fast_copy(self) -> None: + class AddModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + x + + class EmptyTargetedOpsPass(ExportPass): + enable_fast_copy = True + 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"] + + pass_ = self._CountingTargetedPass((exir_ops.edge.aten.mul.Tensor,)) + 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_node_debug_str_is_current_on_fast_and_slow_paths(self) -> None: + class AddThenMulModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x + x) * x + + graph_module = export( + AddThenMulModule(), (torch.randn(2),), strict=True + ).graph_module + self._ensure_tensor_meta(graph_module) + expected = { + node.target: node.format_node() + for node in graph_module.graph.nodes + if node.op == "call_function" + } + + class DebugTrackingPass(self._CountingTargetedPass): + def __init__(self) -> None: + super().__init__((torch.ops.aten.mul.Tensor,)) + self.debug_strings: dict[torch.fx.node.Target, str | None] = {} + + def should_fast_copy_node(self, target: torch.fx.node.Target) -> bool: + self.debug_strings[target] = self.node_debug_str + return super().should_fast_copy_node(target) + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.debug_strings[op] = self.node_debug_str + return super().call_operator(op, args, kwargs, meta) + + pass_ = DebugTrackingPass() + + pass_(graph_module) + + self.assertEqual(pass_.debug_strings, expected) + + 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): + enable_fast_copy = True + 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_convolution_and_linear_targets_are_not_fast_copied(self) -> None: + unsafe_targets = ( + torch.ops.aten.convolution, + torch.ops.aten.convolution.default, + torch.ops.aten.conv1d, + torch.ops.aten.conv1d.default, + torch.ops.aten.conv1d.padding, + torch.ops.aten.conv2d, + torch.ops.aten.conv2d.default, + torch.ops.aten.conv2d.padding, + torch.ops.aten.conv3d, + torch.ops.aten.conv3d.default, + torch.ops.aten.conv3d.padding, + torch.ops.aten.conv_transpose1d, + torch.ops.aten.conv_transpose1d.default, + torch.ops.aten.conv_transpose2d, + torch.ops.aten.conv_transpose2d.input, + torch.ops.aten.conv_transpose3d, + torch.ops.aten.conv_transpose3d.input, + torch.ops.aten.linear, + torch.ops.aten.linear.default, + exir_ops.edge.aten.convolution.default, + exir_ops.edge.aten.conv2d.default, + exir_ops.edge.aten.conv2d.padding, + exir_ops.edge.aten.conv3d.default, + exir_ops.edge.aten.conv3d.padding, + exir_ops.edge.aten.linear.default, + ) + pass_ = ExportPass() + + for target in unsafe_targets: + with self.subTest(target=target): + self.assertFalse(pass_.should_fast_copy_node(target)) + self.assertTrue(pass_.should_fast_copy_node(torch.ops.aten.add.Tensor)) + + def test_packet_target_does_not_match_overload_target(self) -> None: + graph_module = self._raw_add_graph_module() + self._ensure_tensor_meta(graph_module) + + pass_ = self._CountingTargetedPass((torch.ops.aten.add,)) + 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)},) + ) + + pass_ = self._CountingTargetedPass((torch.ops.aten.add.Tensor,)) + 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): + enable_fast_copy = True + 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_tuple_result_keeps_downstream_fast_copy_enabled(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 TupleResultPass(ExportPass): + enable_fast_copy = True + 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 + return ( + self.call_getitem(result, 0, meta), + self.call_getitem(result, 1, meta), + ) + + pass_ = TupleResultPass() + new_graph_module = pass_(graph_module).graph_module + test_input = torch.randn(2, 3) + + new_graph_module.graph.lint() + self.assertEqual(pass_.operator_calls, 1) + torch.testing.assert_close( + new_graph_module(test_input), + graph_module(test_input), + ) + + def test_copied_get_attr_is_reused_by_hot_node_and_calls_on_attr(self) -> None: + value = torch.ones(2) + root = torch.nn.Module() + root.register_buffer("weight", value, persistent=False) + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(2) + weight = graph.get_attr("weight") + weight.meta["val"] = value + cold = graph.call_function(torch.ops.aten.add.Tensor, (x, weight)) + cold.meta["val"] = x.meta["val"] + value + cold.meta["tensor_meta"] = _extract_tensor_metadata(cold.meta["val"]) + hot = graph.call_function(torch.ops.aten.mul.Tensor, (cold, weight)) + hot.meta["val"] = cold.meta["val"] * value + hot.meta["tensor_meta"] = _extract_tensor_metadata(hot.meta["val"]) + graph.output(hot) + graph_module = torch.fx.GraphModule(root, graph) + + class AttrTrackingPass(ExportPass): + enable_fast_copy = True + targeted_ops = (torch.ops.aten.mul.Tensor,) + + def __init__(self) -> None: + super().__init__() + self.attrs: list[ProxyValue] = [] + self.targeted_attr: Any = None + + def on_attr(self, attr: ProxyValue) -> None: + self.attrs.append(attr) + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + if op is torch.ops.aten.mul.Tensor: + self.targeted_attr = args[1] + return super().call_operator(op, args, kwargs, meta) + + pass_ = AttrTrackingPass() + new_graph_module = pass_(graph_module).graph_module + get_attrs = list(new_graph_module.graph.find_nodes(op="get_attr")) + + new_graph_module.graph.lint() + self.assertEqual(len(get_attrs), 2) + self.assertTrue(all(node.target == "weight" for node in get_attrs)) + self.assertEqual(len(pass_.attrs), 2) + self.assertTrue(all(attr.data is value for attr in pass_.attrs)) + self.assertIs(pass_.targeted_attr, value) + self.assertIn("weight", new_graph_module._buffers) + self.assertIn("weight", new_graph_module._non_persistent_buffers_set) + self.assertEqual(len(list(new_graph_module.named_buffers())), 1) + torch.testing.assert_close( + new_graph_module(torch.full((2,), 2.0)), + graph_module(torch.full((2,), 2.0)), + ) + + def test_non_module_dotted_get_attr_fast_copy_fallback_is_atomic(self) -> None: + weight = torch.ones(2) + weight.x = torch.full((2,), 2.0) + root = torch.nn.Module() + root.add_module("w", torch.nn.Module()) + root.w.register_buffer("x", weight.x) + graph = torch.fx.Graph() + wx = graph.get_attr("w.x") + wx.meta["val"] = weight.x + cold_node = graph.call_function(torch.ops.aten.add.Tensor, (wx, wx)) + cold_node.meta["val"] = weight.x + 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) + delattr(graph_module, "w") + graph_module.register_buffer("w", weight) + pass_ = self._CountingTargetedPass() + + new_graph_module = pass_(graph_module).graph_module + + new_graph_module.graph.lint() + self.assertEqual(pass_.operator_calls, 1) + self.assertFalse( + any( + isinstance(child, torch.nn.Module) + for child in new_graph_module.children() + ) + ) + self.assertFalse( + any( + node.op == "get_attr" and len(node.users) == 0 + for node in new_graph_module.graph.nodes + ) + ) + torch.testing.assert_close(new_graph_module(), graph_module()) + + def test_overlapping_get_attr_fast_copy_fallback_is_atomic(self) -> None: + 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") + w.meta["val"] = weight + wx = graph.get_attr("w.x") + wx.meta["val"] = weight.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_ = self._CountingTargetedPass() + + 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_missing_source_val_does_not_disable_downstream_fast_copy(self) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(2) + without_val = graph.call_function(torch.ops.aten.add.Tensor, (x, x)) + copied = graph.call_function(torch.ops.aten.mul.Tensor, (x, x)) + copied.meta["val"] = x.meta["val"] * x.meta["val"] + copied.meta["tensor_meta"] = _extract_tensor_metadata(copied.meta["val"]) + graph.output((without_val, copied)) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + class MissingValPass(self._CountingTargetedPass): + def should_fast_copy_node(self, target: torch.fx.node.Target) -> bool: + return target is not torch.ops.aten.add.Tensor + + pass_ = MissingValPass(()) + new_graph_module = pass_(graph_module).graph_module + + new_graph_module.graph.lint() + self.assertEqual(pass_.operator_calls, 1) + expected = graph_module(torch.full((2,), 3.0)) + actual = new_graph_module(torch.full((2,), 3.0)) + torch.testing.assert_close(actual, expected) + + def test_slow_path_metadata_drift_disables_downstream_fast_copy(self) -> None: + class AddThenMulModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x + x) * x + + graph_module = export( + AddThenMulModule(), (torch.randn(2),), strict=True + ).graph_module + self._ensure_tensor_meta(graph_module) + + class DriftPass(ExportPass): + enable_fast_copy = True + 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 torch.ops.aten.add.Tensor + + def call_operator( + self, + op: Any, + args: tuple[Any, ...], + kwargs: dict[str, Any], + meta: NodeMetadata, + ) -> ProxyValue: + self.operator_calls += 1 + result = super().call_operator(op, args, kwargs, meta) + if op is torch.ops.aten.add.Tensor: + return ProxyValue(result.data.to(torch.float64), result.proxy) + return result + + pass_ = DriftPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 2) + + def test_local_fallback_without_drift_keeps_fast_copy_enabled(self) -> None: + value = torch.ones(2) + value.x = torch.full((2,), 2.0) + root = torch.nn.Module() + root.add_module("w", torch.nn.Module()) + root.w.register_buffer("x", value.x) + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(2) + wx = graph.get_attr("w.x") + wx.meta["val"] = value.x + fallback = graph.call_function(torch.ops.aten.add.Tensor, (wx, wx)) + fallback.meta["val"] = value.x + value.x + fallback.meta["tensor_meta"] = _extract_tensor_metadata(fallback.meta["val"]) + copied = graph.call_function(torch.ops.aten.mul.Tensor, (x, x)) + copied.meta["val"] = x.meta["val"] * x.meta["val"] + copied.meta["tensor_meta"] = _extract_tensor_metadata(copied.meta["val"]) + graph.output((fallback, copied)) + graph_module = torch.fx.GraphModule(root, graph) + delattr(graph_module, "w") + graph_module.register_buffer("w", value) + + pass_ = self._CountingTargetedPass(()) + new_graph_module = pass_(graph_module).graph_module + + new_graph_module.graph.lint() + self.assertEqual(pass_.operator_calls, 1) + expected = graph_module(torch.full((2,), 3.0)) + actual = new_graph_module(torch.full((2,), 3.0)) + torch.testing.assert_close(actual, expected) + + def test_node_backed_proxy_value_supports_fast_copy(self) -> None: + graph_module = self._raw_add_graph_module() + self._ensure_tensor_meta(graph_module) + + class NodeBackedPlaceholderPass(ExportPass): + enable_fast_copy = True + targeted_ops: tuple[()] = () + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def placeholder( + self, name: str, arg: Any, meta: NodeMetadata + ) -> ProxyValue: + result = super().placeholder(name, arg, meta) + return ProxyValue(result.data, result.node) + + 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_ = NodeBackedPlaceholderPass() + new_graph_module = pass_(graph_module).graph_module + + new_graph_module.graph.lint() + self.assertEqual(pass_.operator_calls, 0) + torch.testing.assert_close( + new_graph_module(torch.full((2,), 3.0)), + graph_module(torch.full((2,), 3.0)), + ) + + def test_placeholder_metadata_drift_disables_fast_copy(self) -> None: + graph_module = self._raw_add_graph_module() + self._ensure_tensor_meta(graph_module) + + class PlaceholderDriftPass(ExportPass): + enable_fast_copy = True + targeted_ops: tuple[()] = () + + def __init__(self) -> None: + super().__init__() + self.operator_calls = 0 + + def placeholder( + self, name: str, arg: Any, meta: NodeMetadata + ) -> ProxyValue: + result = super().placeholder(name, arg, meta) + return ProxyValue(result.data.to(torch.float64), result.proxy) + + 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_ = PlaceholderDriftPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 1) + + def test_unknown_metadata_leaf_drift_disables_fast_copy(self) -> None: + graph_module = self._raw_add_graph_module() + self._ensure_tensor_meta(graph_module) + graph = graph_module.graph + output = next(node for node in graph.nodes if node.op == "output") + add = self._single_call_function_node(graph_module, torch.ops.aten.add.Tensor) + + class UnknownMetadataLeaf: + def __repr__(self) -> str: + raise AssertionError("metadata comparison must not call repr") + + add.meta["val"] = (add.meta["val"], UnknownMetadataLeaf()) + with graph.inserting_after(add): + mul = graph.call_function(torch.ops.aten.mul.Tensor, add.args) + mul.meta["val"] = add.meta["val"][0] * add.meta["val"][0] + mul.meta["tensor_meta"] = _extract_tensor_metadata(mul.meta["val"]) + output.args = (mul,) + graph_module.recompile() + + class UnknownLeafDriftPass(ExportPass): + enable_fast_copy = True + 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, + ) -> Any: + self.operator_calls += 1 + result = super().call_operator(op, args, kwargs, meta) + if op is torch.ops.aten.add.Tensor: + return (result, UnknownMetadataLeaf()) + return result + + pass_ = UnknownLeafDriftPass() + + pass_(graph_module) + + self.assertEqual(pass_.operator_calls, 2) + + def test_foreign_node_backed_proxy_value_disables_fast_copy(self) -> None: + graph_module = self._raw_add_graph_module() + self._ensure_tensor_meta(graph_module) + + graph = graph_module.graph + output = next(node for node in graph.nodes if node.op == "output") + add = self._single_call_function_node(graph_module, torch.ops.aten.add.Tensor) + with graph.inserting_after(add): + mul = graph.call_function(torch.ops.aten.mul.Tensor, add.args) + mul.meta.update(add.meta) + output.args = (mul,) + graph_module.recompile() + + class ForeignNodePass(ExportPass): + enable_fast_copy = True + 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, + ) -> Any: + self.operator_calls += 1 + result = super().call_operator(op, args, kwargs, meta) + if op is torch.ops.aten.add.Tensor: + foreign_node = torch.fx.Graph().placeholder("foreign") + return ProxyValue(result.data, foreign_node) + return result + + pass_ = ForeignNodePass() + new_graph_module = pass_(graph_module).graph_module + + new_graph_module.graph.lint() + self.assertEqual(pass_.operator_calls, 2) + + def test_hot_to_cold_dependency_is_remapped_and_executable(self) -> None: + class AddThenMulModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return (x + x) * x + + graph_module = export( + AddThenMulModule(), (torch.randn(2),), strict=True + ).graph_module + self._ensure_tensor_meta(graph_module) + + pass_ = self._CountingTargetedPass((torch.ops.aten.add.Tensor,)) + new_graph_module = pass_(graph_module).graph_module + + new_graph_module.graph.lint() + self.assertEqual(pass_.operator_calls, 1) + torch.testing.assert_close( + new_graph_module(torch.full((2,), 3.0)), + graph_module(torch.full((2,), 3.0)), + ) + + def test_on_attr_runtime_error_propagates(self) -> None: + root = torch.nn.Module() + root.register_buffer("weight", torch.ones(2)) + graph = torch.fx.Graph() + weight = graph.get_attr("weight") + cold = graph.call_function(torch.ops.aten.add.Tensor, (weight, weight)) + cold.meta["val"] = root.weight + root.weight + cold.meta["tensor_meta"] = _extract_tensor_metadata(cold.meta["val"]) + graph.output(cold) + graph_module = torch.fx.GraphModule(root, graph) + + class RaisingOnAttrPass(ExportPass): + enable_fast_copy = True + targeted_ops: tuple[()] = () + + def on_attr(self, _attr: ProxyValue) -> None: + raise RuntimeError("unrelated on_attr failure") + + with self.assertRaisesRegex(RuntimeError, "unrelated on_attr failure"): + RaisingOnAttrPass()(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: """ @@ -324,10 +1150,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # Verify constants were doubled for key, original_const in original_values.items(): new_const = result.exported_program.constants[key] - self.assertTrue( - torch.allclose(new_const, original_const * 2), - f"Constant {key} was not doubled correctly", - ) + torch.testing.assert_close(new_const, original_const * 2) def test_adds_constant_to_exported_program(self) -> None: """ @@ -393,11 +1216,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: # Verify the new constant was added to constants dict self.assertEqual(len(result.exported_program.constants), 1) self.assertIn("_test_added_constant", result.exported_program.constants) - self.assertTrue( - torch.allclose( - result.exported_program.constants["_test_added_constant"], - torch.tensor([1.0, 2.0, 3.0]), - ) + torch.testing.assert_close( + result.exported_program.constants["_test_added_constant"], + torch.tensor([1.0, 2.0, 3.0]), ) # Verify input_specs was updated @@ -413,7 +1234,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: op="placeholder" ) ] - self.assertTrue(len(placeholder_names) == 2) + self.assertEqual(len(placeholder_names), 2) # Verify the new input spec has the correct kind new_spec = None