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..f3fd93678ee 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -7,14 +7,16 @@ # pyre-strict import operator +import threading import traceback from abc import ABC, abstractmethod -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from dataclasses import dataclass from typing import ( Any, Callable, Dict, + Iterator, List, MutableMapping, Optional, @@ -27,6 +29,7 @@ ) import torch +from torch._library import utils as _library_utils from executorch.exir import memory from executorch.exir.delegate import executorch_call_delegate, is_lowered_module from executorch.exir.dialects.edge._ops import EdgeOpOverload @@ -34,7 +37,7 @@ from torch import fx from torch._dispatch.python import enable_python_dispatcher from torch._subclasses import FakeTensorMode, UnsupportedFakeTensorException -from torch._subclasses.fake_tensor import FakeTensor +from torch._subclasses.fake_tensor import _DispatchCacheBypassEntry, FakeTensor from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode from torch.export import ExportedProgram from torch.fx import traceback as fx_traceback @@ -143,6 +146,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 +262,10 @@ class ExportPassBaseError(RuntimeError): pass +class _FastCopyFallback(RuntimeError): + pass + + @dataclass(frozen=True) class ExportedProgramPassResult: exported_program: ExportedProgram @@ -280,6 +318,79 @@ def ensures(self, exported_program: ExportedProgram) -> None: # noqa: B027 """ +_EXTRA_CACHEABLE_NAMESPACES: frozenset[str] = frozenset( + { + "quantized_decomposed", + "tosa", + "cortex_m", + } +) + +_FAKETENSOR_CACHE_PATCH_LOCK = threading.RLock() + + +def _is_extra_cacheable_op(op: object) -> bool: + return ( + isinstance(op, torch._ops.OpOverload) + and op.namespace in _EXTRA_CACHEABLE_NAMESPACES + ) + + +def _evict_extra_nonbuiltin_bypasses( + cache: MutableMapping[object, object], +) -> None: + for key, value in tuple(cache.items()): + flattened_key = getattr(key, "key", ()) + if ( + flattened_key + and _is_extra_cacheable_op(flattened_key[0]) + and isinstance(value, _DispatchCacheBypassEntry) + and value.reason == "non-builtin" + ): + cache.pop(key, None) + + +@contextmanager +def _extend_faketensor_cache_builtins( + fake_tensor_mode: FakeTensorMode, +) -> Iterator[None]: + """Allow vetted ExecuTorch namespaces in the FakeTensor dispatch cache.""" + with _FAKETENSOR_CACHE_PATCH_LOCK: + original_is_builtin = _library_utils.is_builtin + + def extended_is_builtin(op: torch._ops.OpOverload) -> bool: + return _is_extra_cacheable_op(op) or original_is_builtin(op) + + try: + _library_utils.is_builtin = extended_is_builtin # pyre-ignore[8] + _evict_extra_nonbuiltin_bypasses(FakeTensorMode.cache) + if fake_tensor_mode.shape_env is not None: + _evict_extra_nonbuiltin_bypasses( + fake_tensor_mode.shape_env.fake_tensor_cache + ) + # Keep positive entries after restoring the predicate: reuse across + # successive ExportPass instances is the performance benefit. + yield + finally: + _library_utils.is_builtin = original_is_builtin # pyre-ignore[8] + + +_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 +524,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 +673,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 +1072,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(): @@ -818,12 +1126,21 @@ def call(self, graph_module: fx.GraphModule) -> PassResult: self.fake_tensor_mode = fake_tensor_mode with fake_tensor_mode, dispatcher_mode: # type: ignore[assignment, union-attr] - result = self.call_submodule(graph_module, tuple(inputs)) + with _extend_faketensor_cache_builtins(fake_tensor_mode): + result = self.call_submodule(graph_module, tuple(inputs)) return result 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..2d5d890ddfb 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -8,11 +8,13 @@ # pyre-strict import unittest +from typing import Any import executorch.exir as exir import torch from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ( + _extend_faketensor_cache_builtins, ExportedProgramPassBase, ExportedProgramPassResult, ExportPass, @@ -27,6 +29,8 @@ 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._library import utils as _library_utils +from torch._subclasses import FakeTensorMode class TestPassInfra(unittest.TestCase): @@ -228,6 +232,132 @@ 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 TestExportPassFakeTensorCache(unittest.TestCase): + def test_extension_restores_builtin_predicate_after_exception(self) -> None: + op = torch.ops.quantized_decomposed.quantize_per_tensor.default + fake_tensor_mode = FakeTensorMode() + original_is_builtin = _library_utils.is_builtin + + self.assertFalse(original_is_builtin(op)) + with self.assertRaisesRegex(RuntimeError, "test failure"): + with _extend_faketensor_cache_builtins(fake_tensor_mode): + self.assertTrue(_library_utils.is_builtin(op)) + raise RuntimeError("test failure") + + self.assertIs(_library_utils.is_builtin, original_is_builtin) + + def test_extension_replaces_nonbuiltin_bypass_with_cache_entry(self) -> None: + op = torch.ops.quantized_decomposed.quantize_per_tensor.default + fake_tensor_mode = FakeTensorMode() + fake_x = fake_tensor_mode.from_tensor(torch.randn(4)) + FakeTensorMode.cache_clear() + + try: + with fake_tensor_mode: + op(fake_x, 0.1, 0, -128, 127, torch.int8) + bypasses = FakeTensorMode.cache_info().bypasses + self.assertEqual(bypasses.get("non-builtin"), 1) + + with fake_tensor_mode, _extend_faketensor_cache_builtins( + fake_tensor_mode + ): + op(fake_x, 0.1, 0, -128, 127, torch.int8) + after_miss = FakeTensorMode.cache_info() + op(fake_x, 0.1, 0, -128, 127, torch.int8) + after_hit = FakeTensorMode.cache_info() + + self.assertEqual(after_miss.misses, 1) + self.assertEqual(after_hit.hits, after_miss.hits + 1) + finally: + FakeTensorMode.cache_clear() + + class TestExportedProgramPassManager(unittest.TestCase): def test_runs_graph_module_passes_on_exported_program(self) -> None: """