diff --git a/backends/aoti/aoti_partitioner.py b/backends/aoti/aoti_partitioner.py index f84febbdc24..0a1a29b3040 100644 --- a/backends/aoti/aoti_partitioner.py +++ b/backends/aoti/aoti_partitioner.py @@ -100,13 +100,18 @@ def is_node_supported( node.meta["delegation_tag"] = user.meta["delegation_tag"] break - tag_constant_data(exported_program) - tag_mutated_buffer(exported_program) - # tag_constant_data only tags constants that have users; tag the - # genuinely unused ones too so none are left dangling. + # genuinely unused ones too so none are left dangling. A mutated buffer + # gets the tag of the node that writes its new value, so it stays out + # when that value comes from outside the delegates, as in + # self.state.copy_(x). if partition_tags: fallback_tag = next(iter(partition_tags)) + signature = exported_program.graph_signature + writers = { + buffer: output for output, buffer in signature.buffers_to_mutate.items() + } + nodes = {node.name: node for node in exported_program.graph.nodes} for node in exported_program.graph.nodes: if ( node.op == "placeholder" @@ -118,7 +123,17 @@ def is_node_supported( or is_lifted_tensor_constant(exported_program, node) ) ): - node.meta["delegation_tag"] = fallback_tag + writer = writers.get(signature.inputs_to_buffers.get(node.name)) + tag = ( + fallback_tag + if writer is None + else nodes[writer].meta.get("delegation_tag") + ) + if tag is not None: + node.meta["delegation_tag"] = tag + + tag_constant_data(exported_program) + tag_mutated_buffer(exported_program) return PartitionResult( tagged_exported_program=exported_program, partition_tags=partition_tags diff --git a/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in b/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in index 8ce968760aa..231279b6762 100644 --- a/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in +++ b/backends/arm/scripts/docgen/vgf/backends-arm-vgf-overview.md.in @@ -42,7 +42,8 @@ The full user-facing API is documented below. For static mutable buffers such as KV caches, set `alias_buffer_mutations=True`. Eligible buffers must be zero-initialized `int8` tensors whose reads and mutation -are contained in one delegated partition. The backend aliases each buffer's +are contained in one delegated partition, and whose new value is not returned or +read outside that partition. The backend aliases each buffer's input and mutation output to persistent VGF storage across invocations. $COMPILE_SPEC diff --git a/backends/arm/test/misc/test_compile_spec.py b/backends/arm/test/misc/test_compile_spec.py index 7658884cdc0..20b39b13bf5 100644 --- a/backends/arm/test/misc/test_compile_spec.py +++ b/backends/arm/test/misc/test_compile_spec.py @@ -173,6 +173,7 @@ def fake_partition(_self, _exported_program): monkeypatch.setattr(TOSAPartitioner, "partition", fake_partition) monkeypatch.setattr(partitioner, "_validate_mutable_buffers", lambda _: None) monkeypatch.setattr(vgf_partitioner, "tag_mutated_buffer", lambda _: None) + monkeypatch.setattr(partitioner, "_check_mutable_buffers_owned", lambda _: None) result = partitioner.partition(tagged_program) tag0_specs = result.partition_tags["tag0"].compile_specs @@ -193,6 +194,37 @@ def fake_partition(_self, _exported_program): assert all(spec.key != "mutable_buffer_pairs" for spec in shared_spec.compile_specs) +@mark.parametrize("returned", [False, True]) +def test_vgf_aliasing_rejects_a_buffer_the_delegate_cannot_own(monkeypatch, returned): + class MutatingCache(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(4, dtype=torch.int8)) + + def forward(self, x): + self.cache.add_(x) + # A delegate cannot own a buffer whose new value is returned. + return self.cache if returned else self.cache + 1 + + partitioner = VgfPartitioner(VgfCompileSpec(alias_buffer_mutations=True)) + + def tag_every_op(_self, program): + for node in program.graph.nodes: + if node.op == "call_function": + node.meta["delegation_tag"] = "tag0" + return PartitionResult(program, {"tag0": partitioner.delegation_spec}) + + monkeypatch.setattr(TOSAPartitioner, "partition", tag_every_op) + exported_program = torch.export.export( + MutatingCache(), (torch.ones(4, dtype=torch.int8),) + ).run_decompositions() + if returned: + with raises(ValueError, match="cannot be aliased"): + partitioner.partition(exported_program) + else: + partitioner.partition(exported_program) + + def test_vgf_direct_aliasing_requires_partitioner(monkeypatch): class MutatingCache(torch.nn.Module): diff --git a/backends/arm/vgf/partitioner.py b/backends/arm/vgf/partitioner.py index 34dfbacc657..90aa757493d 100644 --- a/backends/arm/vgf/partitioner.py +++ b/backends/arm/vgf/partitioner.py @@ -103,6 +103,22 @@ def _validate_mutable_buffers(self, exported_program: ExportedProgram) -> None: f"mutation tag={mutation_tag}, consumers=[{consumers}]." ) + def _check_mutable_buffers_owned(self, exported_program: ExportedProgram) -> None: + # tag_mutated_buffer also keeps a buffer out of the delegate when its + # new value is returned or read outside it, which + # _validate_mutable_buffers does not see. Aliasing needs the delegate + # to own the buffer. + signature = exported_program.graph_signature + mutated = set(signature.buffers_to_mutate.values()) + for node in exported_program.graph.nodes: + buffer_name = signature.inputs_to_buffers.get(node.name) + if buffer_name in mutated and "delegation_tag" not in node.meta: + raise ValueError( + f"VGF mutable buffer {buffer_name!r} cannot be aliased: its " + "new value is returned or read outside the delegated " + "partition." + ) + def partition(self, exported_program: ExportedProgram) -> PartitionResult: """Partition the program and tag TOSA-compatible subgraphs. @@ -133,4 +149,5 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: } self._validate_mutable_buffers(result.tagged_exported_program) tag_mutated_buffer(result.tagged_exported_program) + self._check_mutable_buffers_owned(result.tagged_exported_program) return result diff --git a/backends/cuda/passes/tests/test_replace_complex_with_real.py b/backends/cuda/passes/tests/test_replace_complex_with_real.py index 793fe16555c..67f760c77ab 100644 --- a/backends/cuda/passes/tests/test_replace_complex_with_real.py +++ b/backends/cuda/passes/tests/test_replace_complex_with_real.py @@ -89,9 +89,8 @@ def forward(self, x, positions, scale): torch.testing.assert_close(program.module()(*inputs), model(*inputs)) def test_tensor_with_two_placeholders(self): - """Partitioning gives a tensor a placeholder per op that reads it when - another backend takes ops between those, and the pass stores the - tensor once.""" + """Partitioning gives a tensor a placeholder per delegate that reads it, + and the pass stores the tensor once.""" class Rotate(torch.nn.Module): def __init__(self): @@ -107,6 +106,9 @@ def forward(self, x, y): model = Rotate() inputs = (torch.randn(3, 2), torch.randn(3, 2)) program = _program_for_cuda_passes(model, inputs) + freqs = next(n for n in program.graph.nodes if n.name == "b_freqs") + for i, user in enumerate(freqs.users): + user.meta["delegation_tag"] = f"tag{i}" self.assertEqual( duplicate_constant_node(program, "b_freqs"), {"b_freqs_copy_0"} ) diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 996ed2a20b8..159fdeeaa1f 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -33,6 +33,7 @@ parse_offgraph_kv_manifest, ring_physical_capacity, ) +from executorch.exir import to_edge from executorch.exir._serialize._cord import FileBackedData from executorch.exir._serialize._named_data_store import NamedDataStore from executorch.exir.backend.backend_details import PreprocessResult @@ -818,6 +819,86 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: f"Constant placeholder {node.name} has tag '{actual_tag}' but expected '{expected_tag}'", ) + def test_unread_mutated_buffer_copied_from_an_input_stays_out(self) -> None: + """ + A mutated buffer that no op reads stays in the owning program when its + new value is a method input, as in self.state.copy_(x). When the only + partition computes the new value, the buffer goes to that delegate. + """ + + class WriteState(torch.nn.Module): + def __init__(self, copy_input: bool) -> None: + super().__init__() + self.copy_input = copy_input + self.register_buffer("state", torch.zeros(2)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.state.copy_(x if self.copy_input else x + x) + return torch.sin(x) + + for copy_input in (False, True): + # to_edge removes the copy, so the new value comes straight from x or x + x. + edge = to_edge(export(WriteState(copy_input), (torch.randn(2),))) + result = CudaPartitioner([]).partition(edge.exported_program()) + program = result.tagged_exported_program + state = next( + n + for n in program.graph.nodes + if n.op == "placeholder" and is_buffer(program, n) + ) + tag = state.meta.get("delegation_tag") + if copy_input: + self.assertIsNone(tag) + else: + self.assertIn(tag, result.partition_tags) + + def test_unread_mutated_buffer_goes_to_the_partition_that_writes_it( + self, + ) -> None: + """ + With several partitions, a mutated buffer that no op reads goes to the + delegate whose op computes its new value, even when that partition is + not the one listed first. + """ + + class WriteState(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("state", torch.zeros(2)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + s = x + x + self.state.copy_(s) + return torch.cos(s * 3) + + program = to_edge(export(WriteState(), (torch.randn(2),))).exported_program() + graph = program.graph_module.graph + nodes = {n.name: n for n in graph.nodes} + add, mul, cos = ( + nodes[name] + for name in ("aten_add_tensor", "aten_mul_tensor", "aten_cos_default") + ) + # A delegate of another backend between mul and cos splits this + # backend's ops into two partitions. + program.graph_module.lowered_module_0 = torch.nn.Module() + with graph.inserting_before(cos): + lowered = graph.get_attr("lowered_module_0") + delegate = graph.call_function(executorch_call_delegate, (lowered, mul)) + output = graph.call_function(operator.getitem, (delegate, 0)) + cos.replace_input_with(mul, output) + graph.lint() + + result = CudaPartitioner([]).partition(program) + state = next( + n + for n in program.graph.nodes + if n.op == "placeholder" and is_buffer(program, n) + ) + self.assertNotEqual( + add.meta["delegation_tag"], next(iter(result.partition_tags)) + ) + self.assertEqual(state.meta.get("delegation_tag"), add.meta["delegation_tag"]) + def test_does_not_retag_already_lowered_delegate(self) -> None: """ A node already lowered by a previous partitioner appears as an diff --git a/backends/mlx/builder/program_builder.py b/backends/mlx/builder/program_builder.py index d5ab4124579..2a1c855171e 100644 --- a/backends/mlx/builder/program_builder.py +++ b/backends/mlx/builder/program_builder.py @@ -59,7 +59,11 @@ Tid, Vid, ) -from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir._serialize._named_data_store import ( + _tensor_to_bytes, + NamedDataStore, +) +from executorch.exir.tensor_layout import TensorLayout from torch.export.exported_program import ExportedProgram from torch.fx.node import Node from torch.utils import _pytree as pytree @@ -316,8 +320,10 @@ def release_placeholder_tensor(self, node: Node) -> None: A repacking handler replaces a raw weight with new constants, leaving the original dead. It would otherwise stay in the ExportedProgram until get_named_data_store() sweeps unused constants -- which only runs after - the whole graph is built, so every raw weight and every repacked weight - are live simultaneously (+16 GB on a 16 GB Q4_K model). + the whole graph is built. Dropping it frees the raw weight only when + nothing else holds it: during lowering, the delegate's program shares + its tensors with the owning program (see + MLXBackend.copy_exported_program_for_preprocess). Only call this when the raw data is fully consumed: the lowering must not reference the placeholder's slot in any emitted instruction, or it @@ -520,21 +526,21 @@ def _make_io_slots(self): # noqa: C901 else: raise NotImplementedError(f"Support for input {arg} is not implemented") - target_to_constant_slot = {} + constant_targets = set() for node in self.ep.graph.nodes: if node.op == "placeholder": if node.users == {}: continue if node.name in constant_tensors: - # A constant shared with other partitions is duplicated once per - # user, so several placeholders here can name the same target. + # Constant data is named by target, so a second slot would get none. target = constant_tensors[node.name] - if target in target_to_constant_slot: - self.set_slot(node, target_to_constant_slot[target]) - else: - target_to_constant_slot[target] = self.make_or_get_slot( - node, id_space=IdSpace.Constant + if target in constant_targets: + raise ValueError( + f"constant '{target}' has more than one placeholder, " + f"'{node.name}' is a second one" ) + constant_targets.add(target) + self.make_or_get_slot(node, id_space=IdSpace.Constant) elif node.name in user_inputs: val = node.meta.get("val", None) if isinstance(val, torch.Tensor) and not val.is_contiguous(): @@ -1132,20 +1138,29 @@ def _build_mlx_graph(self) -> MLXGraph: tensor_meta=tensor_meta_list, ) - def get_named_data_store(self) -> NamedDataStore: + def get_named_data_store( + self, constant_bytes: Dict[int, Tuple[torch.Tensor, bytes]] + ) -> NamedDataStore: """ Get a NamedDataStore containing all constant tensors. + ``constant_bytes`` maps the id of a tensor of the program to the tensor + and its bytes. Delegates that share a constant read the same tensor, so + passing one map to all of them stores its bytes once rather than once + per delegate. The map holds the tensor so that its id stays unique, and + takes only the program's tensors, which outlive the lowering anyway, + not the constants a handler creates. + Uses the unprefixed canonical-name → Slot mapping built by ``_build_io_maps()`` so that tensor lookups hit ``ep.state_dict`` / ``ep.constants`` / ``extra_constants`` (which all use unprefixed keys). The prefix is applied at the exit boundary — the ``NamedDataStore`` key — so it matches the FlatBuffer ``named_slots``. - To reduce peak memory, each constant is deleted from the EP - immediately after its bytes are added to the NamedDataStore. - This avoids holding two full copies of all constants simultaneously - (important for large models where constants can be 20+ GB). + Each constant is deleted from the EP immediately after its bytes are + added to the NamedDataStore. This frees constants a handler created. A + tensor of the program stays alive, held by ``constant_bytes`` (and, + during lowering, by the owning program). """ named_data_store = NamedDataStore() @@ -1155,7 +1170,8 @@ def get_named_data_store(self) -> NamedDataStore: key=lambda x: self._slot_to_final_tid.get(x[1], 0), ) - # Free EP constants not used by the MLX graph to reduce peak memory. + # Drop EP constants not used by the MLX graph; this frees only those + # that nothing else holds. used = set(self._constant_name_to_slot.keys()) for ispec in self.ep.graph_signature.input_specs: if ispec.arg.name in used and ispec.target is not None: @@ -1166,6 +1182,12 @@ def get_named_data_store(self) -> NamedDataStore: if name not in used and isinstance(d[name], torch.Tensor): del d[name] + program_tensors = { + id(tensor) + for d in (self.ep._state_dict, self.ep._constants) + for tensor in d.values() + } + logger.debug(f"Adding {len(entries)} constants to NamedDataStore...") for canonical_name, _slot in entries: tensor = self._find_constant_tensor(canonical_name) @@ -1181,17 +1203,21 @@ def get_named_data_store(self) -> NamedDataStore: ) t = tensor.detach().cpu().contiguous() + data, layout = t, None + if id(tensor) in program_tensors: + if id(tensor) not in constant_bytes: + constant_bytes[id(tensor)] = (tensor, _tensor_to_bytes(t)) + data = constant_bytes[id(tensor)][1] + layout = TensorLayout.from_tensor(t) named_data_store.add_named_data( key=self._prefix_key(canonical_name), - data=t, + data=data, alignment=16, + tensor_layout=layout, ) - # Free the original tensor from the EP immediately. - # The contiguous copy is now serialized as bytes in the - # NamedDataStore — the EP reference is no longer needed. - # (It would be deleted by lowered_backend_module.py after - # preprocess() returns anyway.) + # The bytes are in the NamedDataStore, so the EP reference is no + # longer needed. This frees the tensor only if nothing else holds it. self._delete_constant_tensor(canonical_name) del tensor, t @@ -1231,7 +1257,7 @@ def _find_constant_tensor(self, name: str) -> Optional[torch.Tensor]: return d[k] def _delete_constant_tensor(self, name: str) -> None: - """Delete a constant from the EP to free memory during serialization.""" + """Delete a constant from the EP; this frees it only if nothing else does.""" result = self._resolve_constant(name) if result: diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 5db12feef82..b4704067402 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -2406,6 +2406,53 @@ def _roll_handler(P: MLXProgramBuilder, n: Node) -> Slot: return out +def _constant_mask_as_positions( + P: MLXProgramBuilder, n: Node, non_none: List[Tuple[int, Slot, Any]] +) -> List[Tuple[int, Slot, Any]]: + """Replace a constant boolean index of aten.index.Tensor with its positions. + + The MLX runtime gathers only with integer indices. A boolean mask that is + a constant of the program, used as the only index, becomes one integer + index per mask dimension. Any other boolean index is refused, so the CPU + runs it. + """ + masks = [ + (axis, idx) + for axis, idx, val in non_none + if isinstance(val, torch.Tensor) and val.dtype == torch.bool + ] + if not masks: + return non_none + axis, mask = masks[0] + mask_node = n.args[1][axis] + # The partitioner sees parameters and persistent buffers as fake tensors, + # so only a mask in ep.constants has values it can read. + if ( + len(non_none) > 1 + or mask.id_space != IdSpace.Constant + or (target := P.get_placeholder_target(mask_node)) not in P.ep.constants + ): + raise NotImplementedError( + "aten.index.Tensor: the MLX runtime gathers only with integer " + "indices, so only a constant boolean mask used as the only index " + "is lowered" + ) + value = P.ep.constants[target] + if not value.any(): + raise NotImplementedError( + "aten.index.Tensor: an all-false mask selects nothing, and the MLX " + "runtime cannot write an empty output" + ) + return [ + ( + axis + i, + P.make_or_get_constant(f"{target}_positions_{i}", positions), + positions, + ) + for i, positions in enumerate(value.nonzero(as_tuple=True)) + ] + + @REGISTRY.register(target=[torch.ops.aten.index.Tensor]) def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot: args = P.args(n) @@ -2422,14 +2469,19 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot: x_ndim = len(x_meta.shape) if x_meta is not None else None # Filter out None indices and track which axes they correspond to - non_none = [(i, idx) for i, idx in enumerate(idx_list) if idx is not None] + non_none = [ + (i, idx, n.args[1][i].meta.get("val")) + for i, idx in enumerate(idx_list) + if idx is not None + ] if len(non_none) == 0: raise ValueError("aten.index.Tensor: all indices are None") + non_none = _constant_mask_as_positions(P, n, non_none) + if len(non_none) == 1: - axis, idx = non_none[0] - idx_meta = n.args[1][axis].meta.get("val") + axis, idx, idx_meta = non_none[0] ndim_match = ( x_meta is not None and idx_meta is not None @@ -2464,8 +2516,8 @@ def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot: "aten.index.Tensor with multiple indices requires input shape metadata" ) - indices = [P.slot_to_tid(idx) for _, idx in non_none] - axes = [i for i, _ in non_none] + indices = [P.slot_to_tid(idx) for _, idx, _ in non_none] + axes = [i for i, _, _ in non_none] # slice_sizes: 1 for indexed axes, full static size for non-indexed axes. indexed_axes = set(axes) @@ -2619,6 +2671,10 @@ def _scatter_add_handler(P: MLXProgramBuilder, n: Node) -> Slot: args = P.args(n) require_args(args, 4, 4, "aten.scatter_add") require_kwargs(P.kwargs(n), set(), "aten.scatter_add") + dtype = n.meta["val"].dtype + # MLX's Metal scatter kernels reject 8-byte element types at run time. + if dtype.itemsize == 8: + raise NotImplementedError(f"aten.scatter_add on {dtype} is not supported") x, dim, indices, src = args out = P.make_or_get_slot(n) P.emit( diff --git a/backends/mlx/preprocess.py b/backends/mlx/preprocess.py index b2e26a4d2b2..ac0b583810d 100644 --- a/backends/mlx/preprocess.py +++ b/backends/mlx/preprocess.py @@ -18,9 +18,11 @@ from __future__ import annotations +import copy import hashlib -from typing import ClassVar, final, List +from typing import ClassVar, Dict, final, List, Tuple +import torch from executorch.backends.mlx._logging import logger from executorch.backends.mlx._memprofile import log_footprint, mem_phase from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder @@ -57,6 +59,42 @@ class MLXBackend(BackendDetails): EXPECTED_MAGIC: ClassVar[bytes] = MAGIC EXPECTED_LENGTH: ClassVar[int] = HEADER_LENGTH + @classmethod + def copy_exported_program_for_preprocess( + cls, edge_program: ExportedProgram, compile_specs: List[CompileSpec] + ) -> ExportedProgram: + """Copy the graph but share the constant tensors. + + exir copies the program of every delegate before it preprocesses any of + them, so a deep copy holds all the weights twice, and a constant shared + by several delegates once more for each of them. preprocess only drops + constants from the copy's tables and never writes to a tensor, so + sharing them is safe. + """ + tensors = { + id(tensor): tensor + for values in (edge_program.state_dict, edge_program.constants) + for tensor in values.values() + if isinstance(tensor, torch.Tensor) + } + return copy.deepcopy(edge_program, tensors) + + @classmethod + def preprocess_multimethod( + cls, + edge_programs: Dict[str, List[ExportedProgram]], + compile_specs: Dict[str, List[List[CompileSpec]]], + ) -> Dict[str, List[PreprocessResult]]: + """Preprocess every delegate, sharing the bytes of the constants they share.""" + constant_bytes: Dict[int, Tuple[torch.Tensor, bytes]] = {} + return { + method: [ + cls._preprocess(program, specs, constant_bytes) + for program, specs in zip(programs, compile_specs[method], strict=True) + ] + for method, programs in edge_programs.items() + } + @staticmethod def preprocess( edge_program: ExportedProgram, @@ -73,6 +111,14 @@ def preprocess( PreprocessResult containing the serialized MLX program and data_store_output with constant tensor data. """ + return MLXBackend._preprocess(edge_program, compile_specs, {}) + + @staticmethod + def _preprocess( + edge_program: ExportedProgram, + compile_specs: List[CompileSpec], + constant_bytes: Dict[int, Tuple[torch.Tensor, bytes]], + ) -> PreprocessResult: logger.debug("MLXBackend.preprocess() called") logger.debug(f"Edge program:\n{edge_program}") log_footprint("preprocess entry") @@ -91,7 +137,7 @@ def preprocess( # Get constant data as NamedDataStore (ET will own this data) with mem_phase("get_named_data_store"): - named_data_store = builder.get_named_data_store() + named_data_store = builder.get_named_data_store(constant_bytes) logger.debug( f" named_data_store entries: {len(named_data_store.pte_data)}" diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 5897fe58bee..32720e76a1e 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -20,8 +20,11 @@ from executorch.backends.mlx.passes import get_default_passes from executorch.backends.mlx.test.test_utils import get_mlx_node_counts from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower +from executorch.exir.lowered_backend_module import LoweredBackendModule +from executorch.exir.program import get_fake_program from executorch.runtime import Runtime from torch.export import export +from torch.export.graph_signature import InputKind, InputSpec, TensorArgument class TestMLXPartitionerRejectsToEdge(unittest.TestCase): @@ -79,7 +82,9 @@ def _run(model, inputs, transform_passes=None): counts = get_mlx_node_counts(path) method = Runtime.get().load_program(path).load_method("forward") out = method.execute(list(inputs))[0] - return counts, delegates, (out - ref).abs().max().item() + assert out.shape == ref.shape, (out.shape, ref.shape) + diff = (out - ref).abs() + return counts, delegates, diff.max().item() if diff.numel() else 0.0 class TestMLXPartitionerConstantInputs(unittest.TestCase): @@ -163,6 +168,22 @@ def forward(self, x, y): self.assertGreater(delegates, 0) self.assertLess(err, 1e-6) + def test_int64_scatter_add_stays_off_the_delegate(self): + class ScatterAdd(nn.Module): + def forward(self, base, index, src): + return base.scatter_add(0, index, src) + 1 + + index = torch.tensor([0, 2, 2, 7]) + for dtype, scatters in ((torch.float32, 1), (torch.int64, 0)): + with self.subTest(dtype=dtype): + counts, delegates, err = _run( + ScatterAdd(), + (torch.zeros(8, dtype=dtype), index, torch.ones(4, dtype=dtype)), + ) + self.assertEqual(counts.get("ScatterAddNode", 0), scatters) + self.assertGreater(delegates, 0) + self.assertEqual(err, 0) + class Sdpa(nn.Module): def __init__(self, is_causal: bool = False): @@ -353,6 +374,66 @@ def forward(self, x): self.assertGreater(delegates, 1) self.assertLess(err, 1e-5) + def test_second_placeholder_for_one_constant_raises(self): + class TwoReads(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("c", torch.randn(4)) + + def forward(self, x): + return x + self.c, x * self.c + + ep = export(TwoReads(), (torch.randn(4),), strict=False) + original = next(n for n in ep.graph.nodes if n.name == "b_c") + with ep.graph.inserting_after(original): + second = ep.graph.placeholder("b_c_second") + second.meta = dict(original.meta) + mul = next(n for n in ep.graph.nodes if n.target == torch.ops.aten.mul.Tensor) + mul.replace_input_with(original, second) + ep.graph_signature.input_specs.insert( + 1, InputSpec(InputKind.BUFFER, TensorArgument("b_c_second"), "c", True) + ) + + for check_only in (True, False): + with self.subTest(check_only=check_only): + builder = MLXProgramBuilder(ep) + with self.assertRaisesRegex(ValueError, "more than one placeholder"): + if check_only: + builder.check_support_only() + else: + builder.build() + + def test_delegates_share_the_bytes_of_a_shared_weight(self): + class SharedLinear(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(4, 4) + + def forward(self, x): + # elu is not lowered, so the second use lands in another delegate. + return self.linear(torch.nn.functional.elu(self.linear(x))) + + model = SharedLinear().eval() + weight = model.linear.weight.detach().clone() + program = to_edge_transform_and_lower( + export(model, (torch.randn(2, 4),), strict=False), + partitioner=[MLXPartitioner()], + ).exported_program() + delegates = [ + module + for module in program.graph_module.children() + if isinstance(module, LoweredBackendModule) + ] + + self.assertEqual(len(delegates), 2) + first, second = ( + {id(buffer) for buffer in delegate.named_data_store_output.buffers} + for delegate in delegates + ) + self.assertEqual(len(first), 2) + self.assertEqual(first, second) + self.assertTrue(torch.equal(model.linear.weight, weight)) + class TestMLXPartitionerGroupedKeys(unittest.TestCase): """Repeated KV heads are fused only when the caller enables the shared passes.""" @@ -443,6 +524,126 @@ def forward(self, x): self.assertLess(err, 1e-6) +class TestMLXPartitionerBoolIndex(unittest.TestCase): + """The MLX runtime gathers only with integer indices.""" + + class ConstantMask(nn.Module): + def __init__(self, mask, persistent=None, axis=0, scale=True): + super().__init__() + if persistent is None: + self.mask = mask + else: + self.register_buffer("mask", mask, persistent=persistent) + self.axis = axis + self.scale = scale + + def forward(self, x): + y = (x + 1)[(slice(None),) * self.axis + (self.mask,)] + if self.axis: + # The leading dim's stride depends on how many elements the + # mask selects, which exir spec propagation rejects, so keep it + # in the delegate. + y = y.sum(self.axis) + return y * 2 if self.scale else y + + def test_constant_mask_becomes_positions(self): + mask = torch.tensor([1, 0, 1, 1, 0, 1, 0, 1]) == 1 + cases = { + "1d": (self.ConstantMask(mask), (8, 16)), + "2d": (self.ConstantMask(mask.reshape(1, 8)), (1, 8, 16)), + "3d": (self.ConstantMask(mask.reshape(2, 2, 2)), (2, 2, 2, 5)), + "buffer": (self.ConstantMask(mask, persistent=False), (8, 16)), + "axis_1": (self.ConstantMask(mask.reshape(2, 4), axis=1), (3, 2, 4, 5)), + } + for name, (model, shape) in cases.items(): + with self.subTest(name): + _, delegates, err = _run(model.eval(), (torch.randn(shape),)) + self.assertEqual(delegates, 1) + self.assertLess(err, 1e-5) + + def test_other_bool_indices_run_on_the_cpu(self): + class ComputedMask(nn.Module): + def __init__(self): + super().__init__() + self.ids = torch.tensor([[1, 7, 7, 2, 7, 3, 7, 4]]) + + def forward(self, x): + return (x + 1)[self.ids == 7] * 2 + + class MaskInput(nn.Module): + def forward(self, x, mask): + return (x + 1)[mask] * 2 + + class MaskWithIndex(nn.Module): + def __init__(self, mask): + super().__init__() + self.mask = mask + + def forward(self, x, i): + return (x + 1)[self.mask, i] * 2 + + class MaskThenUpdate(nn.Module): + def __init__(self, mask): + super().__init__() + self.register_buffer("mask", mask, persistent=False) + + def forward(self, x): + y = (x + 1)[self.mask] * 2 + self.mask.copy_(x[:, :, 0] > 0) + return y + + mask = torch.tensor([[1, 0, 1, 1, 0, 1, 0, 1]]) == 1 + x = torch.randn(1, 8, 16) + # An all-false mask gives an empty result, so it is the graph output + # here: an MLX op after it would have to write an empty output too. + all_false = self.ConstantMask(torch.zeros(1, 8, dtype=torch.bool), scale=False) + cases = { + "computed_mask": (ComputedMask(), (x,), 2), + "mask_input": (MaskInput(), (x, mask), 2), + "mask_with_index": (MaskWithIndex(mask), (x, torch.tensor([3])), 2), + "persistent_buffer": (self.ConstantMask(mask, persistent=True), (x,), 2), + "all_false": (all_false, (x,), 1), + } + for name, (model, inputs, expected) in cases.items(): + with self.subTest(name): + _, delegates, err = _run(model.eval(), inputs) + self.assertEqual(delegates, expected) + self.assertLess(err, 1e-6) + + with self.subTest("mutated_buffer"): + # Not run: the eager reference call would overwrite the mask first. + program = _lower(MaskThenUpdate(mask).eval(), (x,)) + self.assertEqual(_delegate_count(program), 2) + + def test_handler_refuses_masks_it_cannot_lower(self): + mask = torch.tensor([1, 0, 1, 1, 0, 1, 0, 1]) == 1 + cases = { + # A persistent buffer has no values in the fake copy, so the handler + # must refuse it rather than fail reading it. + "persistent_buffer": ( + self.ConstantMask(mask, persistent=True), + "only a constant boolean mask", + ), + "all_false": ( + self.ConstantMask(torch.zeros(8, dtype=torch.bool)), + "all-false", + ), + } + for name, (model, reason) in cases.items(): + with self.subTest(name): + # The same fake copy of the program that to_backend gives the + # partitioner. + ep = get_fake_program(export(model.eval(), (torch.randn(8, 16),))) + index = next( + n for n in ep.graph.nodes if n.target == torch.ops.aten.index.Tensor + ) + builder = MLXProgramBuilder(ep) + builder.check_support_only() + info = builder.node_info[index] + self.assertFalse(info.supported) + self.assertIn(reason, info.unsupported_reason) + + class TestMLXPartitionerMixedSupport(unittest.TestCase): """An operator is preserved from decomposition per operator, not per call.""" diff --git a/backends/transforms/test/test_create_delete_constant_placeholder.py b/backends/transforms/test/test_create_delete_constant_placeholder.py index b85a14b036f..d805b4d23f5 100644 --- a/backends/transforms/test/test_create_delete_constant_placeholder.py +++ b/backends/transforms/test/test_create_delete_constant_placeholder.py @@ -211,6 +211,10 @@ def forward(self, x): export(SharedWeight(), (torch.zeros(1),), strict=True) ).exported_program() (weight_name,) = exported_program.graph_signature.inputs_to_parameters + weight = next(n for n in exported_program.graph.nodes if n.name == weight_name) + # A copy is made for a user in another delegate. + for i, user in enumerate(weight.users): + user.meta["delegation_tag"] = f"tag{i}" duplicate_constant_node(exported_program, weight_name) weight_nodes = [ diff --git a/docs/source/backends/arm-vgf/arm-vgf-overview.md b/docs/source/backends/arm-vgf/arm-vgf-overview.md index 2407074b099..142e8d05790 100644 --- a/docs/source/backends/arm-vgf/arm-vgf-overview.md +++ b/docs/source/backends/arm-vgf/arm-vgf-overview.md @@ -42,7 +42,8 @@ The full user-facing API is documented below. For static mutable buffers such as KV caches, set `alias_buffer_mutations=True`. Eligible buffers must be zero-initialized `int8` tensors whose reads and mutation -are contained in one delegated partition. The backend aliases each buffer's +are contained in one delegated partition, and whose new value is not returned or +read outside that partition. The backend aliases each buffer's input and mutation output to persistent VGF storage across invocations. ```python diff --git a/docs/source/backends/cuda/cuda-overview.md b/docs/source/backends/cuda/cuda-overview.md index 56dccf5813b..3f18d5cde6c 100644 --- a/docs/source/backends/cuda/cuda-overview.md +++ b/docs/source/backends/cuda/cuda-overview.md @@ -256,11 +256,14 @@ the delegate, so that the delegate reads the constant itself. `share_mutable_buf planned tensor in host memory, so it does not work with this flag. This flag works on its own: it needs neither `skip_h2d_for_method_inputs` nor unplanned inputs. -**`skip_h2d_for_method_inputs` refuses inputs that the method changes.** `to_executorch` fails if -the method changes an input, even when the delegate makes the change. This includes a mutable -buffer, such as a KV cache, even though the caller does not pass it. For a buffer, put it and every -operator that uses it in the delegate. The CUDA partitioner already moves the buffer in, with -`tag_mutated_buffer`, once those operators are there. +**`skip_h2d_for_method_inputs` refuses changed inputs that a delegate reads.** `to_executorch` fails +if the method changes an input that a delegate reads, even when the delegate makes the change. This +includes a mutable buffer, such as a KV cache, even though the caller does not pass it. For a +buffer, put it and every operator that uses it in the delegate. The CUDA partitioner already moves +the buffer in, with `tag_mutated_buffer`, once those operators are there. It keeps the buffer out +when the method also returns the buffer's new value, or when the new value is a method input, as in +`self.state.copy_(x)`. For a buffer that stays out, this flag refuses the method only if a delegate +reads it. **The choice is baked into the `.pte`.** A program exported with copies expects host tensors; one exported without them expects device tensors. Passing the wrong kind is a caller error, not diff --git a/exir/backend/backend_api.py b/exir/backend/backend_api.py index 8104595bfe1..bc59cb1434a 100644 --- a/exir/backend/backend_api.py +++ b/exir/backend/backend_api.py @@ -304,7 +304,7 @@ def _partition_and_lower_one_graph_module( # Copy the output node meta from the original output node, because # create_submodule_from_nodes doesn't cover the meta field submodule_output_node.meta = tagged_graph_module_output_node.meta - logging.debug(f"Partitioned graph module: {tagged_graph_module}") + logging.debug("Partitioned graph module: %s", tagged_graph_module.graph) ( submodule_program, @@ -423,8 +423,9 @@ def to_backend( update_to_real_program(tagged_exported_program, edge_program) - for tag, _ in partitioner_result.partition_tags.items(): - _maybe_duplicate_constant_nodes(tagged_exported_program, tag) + _maybe_duplicate_constant_nodes( + tagged_exported_program, partitioner_result.partition_tags.keys() + ) tagged_graph_module = _partition_and_lower( tagged_exported_program.graph_module, @@ -486,7 +487,7 @@ def _create_partitions_in_graph_module( submodule_output_node = submodule.graph.output_node() # Copy the output node meta from the original output node, because # create_submodule_from_nodes doesn't cover the meta field - logging.debug(f"Partitioned graph module: {tagged_graph_module}") + logging.debug("Partitioned graph module: %s", tagged_graph_module.graph) ( submodule_program, toplevel_input_specs_to_delete, @@ -750,8 +751,9 @@ def to_backend( update_to_real_program(tagged_exported_program, edge_program) - for tag, _ in partitioner_result.partition_tags.items(): - _maybe_duplicate_constant_nodes(tagged_exported_program, tag) + _maybe_duplicate_constant_nodes( + tagged_exported_program, partitioner_result.partition_tags.keys() + ) backend_id_to_call_submodule_nodes = _create_partitions( tagged_exported_program.graph_module, diff --git a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py index 478c2f35eea..a6d66e3880c 100644 --- a/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py +++ b/exir/backend/canonical_partitioners/duplicate_constant_node_pass.py @@ -6,31 +6,32 @@ import copy import logging -from typing import Optional, Set +from typing import Dict, List, Optional, Set import torch -from torch._export.utils import get_buffer, get_lifted_tensor_constant, get_param from torch.export import ExportedProgram -from torch.export.exported_program import InputSpec, TensorArgument +from torch.export.exported_program import InputKind, InputSpec, TensorArgument -def _get_attribute_or_constants( - exported_program: ExportedProgram, node: torch.fx.Node -) -> Optional[torch.Tensor]: - # get either attribute node or constant constant - maybe_param = get_param(exported_program, node) - maybe_buffer = get_buffer(exported_program, node) - maybe_lifted_tensor = get_lifted_tensor_constant(exported_program, node) - - constant_or_attribute = None - if maybe_param is not None: - constant_or_attribute = maybe_param - elif maybe_buffer is not None: - constant_or_attribute = maybe_buffer - elif maybe_lifted_tensor is not None: - constant_or_attribute = maybe_lifted_tensor - return constant_or_attribute +def _users_reading_a_copy( + node: torch.fx.Node, order: Dict[torch.fx.Node, int] +) -> List[List[torch.fx.Node]]: + """ + The users of node that read a copy of it, in groups that share one. The + users with the same delegation tag, or with none, form a group, so each + program takes a constant once: for example, a program that lists the + same lifted constant twice fails to unlift. The group with the graph + output keeps the original, which the output specs name, and otherwise + the group of the first user in graph order does. + """ + groups: Dict[Optional[str], List[torch.fx.Node]] = {} + # Rewiring a user for another constant moves it to the end of node.users. + for user in sorted(node.users, key=order.__getitem__): + groups.setdefault(user.meta.get("delegation_tag"), []).append(user) + output = any(user.op == "output" for user in node.users) + keeper = None if output else next(iter(groups), None) + return [users for tag, users in groups.items() if tag != keeper] # TODO: add other passes to duplicate call_function nodes @@ -41,6 +42,8 @@ def duplicate_constant_node( A pass to duplicate the attributes/constants node (the candidate_node) in the graph. Mostly used for duplicating light-weight data. If the data is too large, try tag it with "no_copy" to prevent high memory usage and make it as part of the output. A mutated buffer is never copied, because a copy would not see its updates: this raises instead. + Users with the same delegation tag, or with none, share one placeholder, so tag the + users before calling it: on a graph without tags it makes no copies. Args: exported_program: the exported program to be modified. If constants nodes are copied, they will be added as new @@ -50,17 +53,30 @@ def duplicate_constant_node( Returns: The set of the names of the new constant nodes """ - to_be_copied = [ - node + return duplicate_constant_nodes(exported_program, {candidate_node}) + + +def duplicate_constant_nodes( + exported_program: ExportedProgram, candidate_nodes: Set[str] +) -> Set[str]: + """ + duplicate_constant_node for several constant nodes at once. It recompiles + and validates the program once, rather than once per node, which matters + on a large graph with many shared constants. + + Returns: + The set of the names of the new constant nodes + """ + if not any( + node.op == "placeholder" and node.name in candidate_nodes for node in exported_program.graph.nodes - if node.name == candidate_node and node.op == "placeholder" - ] - if len(to_be_copied) == 0: + ): logging.info("no constant node to be copied") return set() new_input_specs = [] old_signature = exported_program.graph_signature copied_nodes = set() + order = {node: i for i, node in enumerate(exported_program.graph.nodes)} placeholder_idx = -1 for node in exported_program.graph.nodes: @@ -70,10 +86,15 @@ def duplicate_constant_node( placeholder_idx += 1 old_input_spec = old_signature.input_specs[placeholder_idx] old_input_spec_copy = copy.deepcopy(old_input_spec) - if node == to_be_copied[0]: + if node.name in candidate_nodes: constant_or_attribute_node = node - constant_or_attribute = _get_attribute_or_constants(exported_program, node) - if constant_or_attribute is None: + # The input spec kind answers this without get_param, get_buffer and + # get_lifted_tensor_constant, which rebuild a dict on every call. + if old_input_spec.kind not in ( + InputKind.PARAMETER, + InputKind.BUFFER, + InputKind.CONSTANT_TENSOR, + ): raise RuntimeError( f"placeholder node for non-params, non-buffer, and non-tensor constants should not be tagged: {node} " ) @@ -87,12 +108,7 @@ def duplicate_constant_node( "that reads or writes it has the same tag, as " "tag_mutated_buffer does, or leave it untagged." ) - # The graph output keeps the original, which the output specs - # name, so then every other user reads a copy. - users = [user for user in node.users if user.op != "output"] - if len(users) == len(node.users): - users = users[1:] - for ith, user in enumerate(users): + for ith, group in enumerate(_users_reading_a_copy(node, order)): copy_constant_or_attribute_fqn = node.name + f"_copy_{ith}" with exported_program.graph.inserting_before( constant_or_attribute_node @@ -113,10 +129,11 @@ def duplicate_constant_node( ) # This also reaches a constant that an op such as cat or # index_put reads from a list. - user.replace_input_with( - constant_or_attribute_node, - copied_constant_or_attribute_node, - ) + for user in group: + user.replace_input_with( + constant_or_attribute_node, + copied_constant_or_attribute_node, + ) # The copy keeps the original target, so it reads the # original tensor and needs no data of its own. new_input_specs.append( diff --git a/exir/backend/test/test_passes.py b/exir/backend/test/test_passes.py index 697f8c78d74..9a66ba3c298 100644 --- a/exir/backend/test/test_passes.py +++ b/exir/backend/test/test_passes.py @@ -34,12 +34,22 @@ def forward(self, x): ) const_nodes = [ - node.name + node for node in edge.exported_program().graph.nodes if node.op == "placeholder" and is_buffer(edge.exported_program(), node) ] + # Users in one program share the original. + self.assertEqual( + duplicate_constant_node(edge.exported_program(), const_nodes[0].name), + set(), + ) + # A copy is made for a user in another delegate. + for i, user in enumerate(const_nodes[0].users): + user.meta["delegation_tag"] = f"tag{i}" - copied_nodes = duplicate_constant_node(edge.exported_program(), const_nodes[0]) + copied_nodes = duplicate_constant_node( + edge.exported_program(), const_nodes[0].name + ) self.assertEqual(len(copied_nodes), 1) # Check that the new constant node is in the graph @@ -76,6 +86,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: user = users[1] user.args = tuple(a for a in user.args if a is not buffer_node) user.kwargs = {"other": buffer_node} + # A copy is made for a user in another delegate. + for i, u in enumerate(users): + u.meta["delegation_tag"] = f"tag{i}" # Patch validation since we modified the graph edge.exported_program()._validate = lambda: None diff --git a/exir/backend/test/test_utils.py b/exir/backend/test/test_utils.py index 24b72e9a1e4..a3b38bef16c 100644 --- a/exir/backend/test/test_utils.py +++ b/exir/backend/test/test_utils.py @@ -4,7 +4,9 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import itertools import unittest +from unittest.mock import patch import torch from executorch import exir @@ -24,6 +26,7 @@ ) from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo from executorch.exir.backend.utils import ( + _maybe_duplicate_constant_nodes, format_delegated_graph, get_delegated_payload, get_delegates, @@ -34,6 +37,7 @@ ) from executorch.exir.dialects._ops import bind_pattern_to_op, ops as exir_ops +from executorch.exir.lowered_backend_module import get_lowered_backend_modules from executorch.extension.pybindings.portable_lib import ( # @manual _get_registered_backend_names, _load_for_executorch_from_buffer, @@ -559,3 +563,170 @@ def partition(self, program: ExportedProgram) -> PartitionResult: program.graph.output_node().args[0][-1].name, "b_const" ) lowered.to_executorch() + + def test_to_backend_gives_each_program_one_copy_of_a_shared_constant(self): + class UseConst(torch.nn.Module): + def __init__(self, cpu_first, kind="constant"): + super().__init__() + const = torch.ones(2) + if kind == "buffer": + self.register_buffer("const", const) + else: + # A plain tensor attribute is lifted as a constant tensor. + self.const = ( + torch.nn.Parameter(const) if kind == "parameter" else const + ) + self.cpu_first = cpu_first + + def forward(self, x): + if self.cpu_first: + s, c = torch.sin(self.const), torch.cos(self.const) + y = (x + self.const) * self.const + z = (x - self.const) / self.const + if not self.cpu_first: + s, c = torch.sin(self.const), torch.cos(self.const) + return y, z, s, c + + class TagNodes(Partitioner): + def __init__(self, tags): + super().__init__() + self.tags = tags + + def partition(self, program: ExportedProgram) -> PartitionResult: + for node in program.graph.nodes: + if node.name in self.tags: + node.meta["delegation_tag"] = self.tags[node.name] + tag_constant_data(program) + return PartitionResult( + tagged_exported_program=program, + partition_tags=dict.fromkeys( + self.tags.values(), + DelegationSpec(ExecutorBackend.__name__, []), + ), + ) + + def counts(program): + programs = [program] + [ + module.original_module + for module in get_lowered_backend_modules(program.graph_module) + ] + return [ + [spec.target for spec in p.graph_signature.input_specs].count("const") + for p in programs + ] + + # Each delegate reads the constant twice, and sin and cos read it + # outside of them. Each program takes it once, whichever reads first. + for kind, cpu_first in itertools.product( + ("constant", "buffer", "parameter"), (False, True) + ): + with self.subTest(kind=kind, cpu_first=cpu_first): + lowered = to_edge( + export(UseConst(cpu_first, kind), (torch.ones(2),), strict=True) + ) + lowered = lowered.to_backend( + TagNodes({"aten_add_tensor": "tag0", "aten_mul_tensor": "tag0"}) + ) + self.assertEqual(counts(lowered.exported_program()), [1, 1]) + lowered.exported_program().module() + lowered = lowered.to_backend( + TagNodes({"aten_sub_tensor": "tag1", "aten_div_tensor": "tag1"}) + ) + self.assertEqual(counts(lowered.exported_program()), [1, 1, 1]) + lowered.to_executorch() + + # One partitioner takes both delegates, so one call copies it twice. + with self.subTest(one_partitioner=True): + lowered = to_edge(export(UseConst(False), (torch.ones(2),), strict=True)) + lowered = lowered.to_backend( + TagNodes( + { + "aten_add_tensor": "tag0", + "aten_mul_tensor": "tag0", + "aten_sub_tensor": "tag1", + "aten_div_tensor": "tag1", + } + ) + ) + self.assertEqual(counts(lowered.exported_program()), [1, 1, 1]) + lowered.to_executorch() + + def test_maybe_duplicate_constant_nodes_refuses_one_tag_string(self): + program = to_edge( + export(torch.nn.Linear(2, 2), (torch.ones(2),), strict=True) + ).exported_program() + # A str is an iterable of tags too, one per character. + with self.assertRaises(TypeError): + _maybe_duplicate_constant_nodes(program, "tag0") + + def test_constants_shared_by_many_delegates_are_copied_in_one_pass(self): + class ShareConsts(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("a", torch.ones(2)) + self.register_buffer("b", torch.ones(2)) + + def forward(self, x): + return x + self.a, x * self.a, self.a - self.b, x / self.b + + program = to_edge( + export(ShareConsts(), (torch.ones(2),), strict=True) + ).exported_program() + ops = [node for node in program.graph.nodes if node.op == "call_function"] + for i, node in enumerate(ops): + node.meta["delegation_tag"] = f"tag{i}" + tag_constant_data(program) + + # Recompiling and validating once per shared constant is quadratic in + # the graph size, which a large model with many delegates cannot afford. + with patch.object( + ExportedProgram, + "_validate", + autospec=True, + side_effect=ExportedProgram._validate, + ) as validate, patch.object( + torch.fx.GraphModule, + "recompile", + autospec=True, + side_effect=torch.fx.GraphModule.recompile, + ) as recompile: + _maybe_duplicate_constant_nodes( + program, [node.meta["delegation_tag"] for node in ops] + ) + self.assertEqual(validate.call_count, 1) + self.assertEqual(recompile.call_count, 1) + + # Copying b_a rewires sub, which must not change who keeps b_b: the + # first reader in graph order does. + self.assertEqual( + { + node.meta["delegation_tag"]: [arg.name for arg in node.args] + for node in ops + }, + { + "tag0": ["x", "b_a"], + "tag1": ["x", "b_a_copy_0"], + "tag2": ["b_a_copy_1", "b_b"], + "tag3": ["x", "b_b_copy_0"], + }, + ) + for node in program.graph.nodes: + if node.op == "placeholder" and node.name != "x": + self.assertEqual( + {user.meta["delegation_tag"] for user in node.users}, + {node.meta["delegation_tag"]}, + ) + self.assertEqual( + { + spec.arg.name: spec.target + for spec in program.graph_signature.input_specs + }, + { + "b_a_copy_0": "a", + "b_a_copy_1": "a", + "b_a": "a", + "b_b_copy_0": "b", + "b_b": "b", + "x": None, + }, + ) diff --git a/exir/backend/utils.py b/exir/backend/utils.py index 96ff686cc2a..bc05f6e4037 100644 --- a/exir/backend/utils.py +++ b/exir/backend/utils.py @@ -16,7 +16,7 @@ import torch from executorch.exir.backend.backend_details import ExportedProgram from executorch.exir.backend.canonical_partitioners.duplicate_constant_node_pass import ( - duplicate_constant_node, + duplicate_constant_nodes, ) from executorch.exir.common import setting_python_recursive_limit from executorch.exir.delegate import executorch_call_delegate @@ -209,7 +209,7 @@ def _assign_new_tag( def _maybe_duplicate_constant_nodes( tagged_exported_program: ExportedProgram, - tag: str, + tags: Iterable[str], ) -> None: """ If the constants node is shared by different tagged nodes, like @@ -222,9 +222,13 @@ def _maybe_duplicate_constant_nodes( backend can estimate how much they want to duplicate the constant node, either error out or default to duplicate """ + if isinstance(tags, str): + raise TypeError("tags is an iterable of tags, not one tag") + tags = set(tags) candidate_nodes = set() for node in tagged_exported_program.graph.nodes: - if node.meta.get("delegation_tag", "") == tag: + tag = node.meta.get("delegation_tag", None) + if tag in tags: if node.op == "placeholder": for user in node.users: users_tag = user.meta.get("delegation_tag", None) @@ -236,12 +240,10 @@ def _maybe_duplicate_constant_nodes( ) else: candidate_nodes.add(node.name) - copied_nodes = set() - for candidate_node in candidate_nodes: - # Both tagged exported program and the owning program need to go through the same duplication pass - copied_nodes = copied_nodes.union( - duplicate_constant_node(tagged_exported_program, candidate_node) - ) + if not candidate_nodes: + return + # Both tagged exported program and the owning program need to go through the same duplication pass + copied_nodes = duplicate_constant_nodes(tagged_exported_program, candidate_nodes) candidate_node_with_copies = candidate_nodes.union(copied_nodes) _assign_new_tag(tagged_exported_program, candidate_node_with_copies) diff --git a/exir/emit/test/test_emit.py b/exir/emit/test/test_emit.py index 6ca6aee763c..cdc8d002696 100644 --- a/exir/emit/test/test_emit.py +++ b/exir/emit/test/test_emit.py @@ -2676,10 +2676,25 @@ def forward(self, x): # it contiguous. return x * self.scale.expand(1) - exported = constant_prop_pass( - export(ScaleByBuffer(), (torch.randn(4),), strict=True) - ) - self.assertEqual([t.stride() for t in exported.constants.values()], [(1,)]) + class AddStridedWeight(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", torch.randn(48)) + + def forward(self, x): + # Channels-last but for a zero stride on the size-1 batch dim. + # The folded constant is contiguous, like the input it is added + # to, so portable add sees one dim order. + return x + self.weight.as_strided((1, 3, 4, 4), (0, 1, 12, 3)) + + for model, x, stride in ( + (ScaleByBuffer(), torch.randn(4), (1,)), + (AddStridedWeight(), torch.randn(1, 3, 4, 4), (48, 16, 4, 1)), + ): + exported = constant_prop_pass(export(model, (x,), strict=True)) + self.assertEqual( + [t.stride() for t in exported.constants.values()], [stride] + ) def test_emit_user_constants_with_zero_stride(self) -> None: def zero_stride(value): diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index ba8f0d6fdb8..68316fe3630 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -7,6 +7,7 @@ # pyre-strict import copy +import heapq import operator from collections import defaultdict from typing import Any, Dict, List, Optional, Set, Tuple, Union @@ -43,7 +44,6 @@ OutputSpec, TensorArgument, ) -from torch.fx.passes.tools_common import stable_topological_sort from torch.fx.passes.utils.fuser_utils import ( erase_nodes, fuse_as_graphmodule, @@ -866,6 +866,35 @@ def create_exported_program_from_submodule( ) +def _stable_topological_sort_in_place(graph: torch.fx.Graph) -> None: + """ + Orders the nodes like torch's stable_topological_sort, by moving them. That + function installs a new graph, and a GraphModule then recompiles itself and + every submodule, so sorting once per partition took time quadratic in the + number of partitions. + """ + position = {node: i for i, node in enumerate(graph.nodes)} + unsorted_inputs = {node: len(node.all_input_nodes) for node in graph.nodes} + ready = [(i, node) for node, i in position.items() if unsorted_inputs[node] == 0] + previous = None + while ready: + _, node = heapq.heappop(ready) + if previous is None: + first = next(iter(graph.nodes)) + if first is not node: + first.prepend(node) + elif previous.next is not node: + previous.append(node) + previous = node + del position[node] + for user in node.users: + unsorted_inputs[user] -= 1 + if unsorted_inputs[user] == 0: + heapq.heappush(ready, (position[user], user)) + if position: + raise AssertionError(f"Input graph has cycles, unable to add {list(position)}") + + def create_submodule_from_nodes( gm: torch.fx.GraphModule, node_list: NodeList, @@ -875,7 +904,8 @@ def create_submodule_from_nodes( """ Modifies the given graph module in-place to separate out the given nodes into a submodule. The given node_list should form a fully connected - subgraph. + subgraph. Unless skip_legalize_graph is set, the nodes of gm are sorted + in place. gm is not recompiled, so call gm.recompile() before running it. Args: gm: The graph module that we want to partition @@ -893,6 +923,7 @@ def create_submodule_from_nodes( ) _fixup_output_node(sub_gm) + sub_gm.recompile() gm = insert_subgm(gm, sub_gm, orig_inputs, orig_outputs) submodule_node = None @@ -934,7 +965,7 @@ def create_submodule_from_nodes( # TODO : T153794167 Get rid of support for skipping the sort in create_submodule_from_nodes # once we transition to using fuse_by_partitions. if not skip_legalize_graph: - stable_topological_sort(gm) + _stable_topological_sort_in_place(gm.graph) # Get the call_module node submodule_node = None diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index 769a9a9c2c1..aa8e4318954 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -20,7 +20,6 @@ _schema_paired_args, ) from executorch.exir.passes.replace_aten_with_edge_pass import aten_to_edge -from executorch.exir.tensor import without_zero_strides from torch._export.utils import ( get_buffer, get_lifted_tensor_constant, @@ -289,13 +288,22 @@ def get_propagated_const_tensor_dict( # Execute the `node.target` and create a new propagated constant tensor. prop_constant_tensor = node.target(*args_data, **kwargs_data) + # ExecuTorch doesn't support zero strides, so we need to ensure the tensor is contiguous + # if it has any zero strides from broadcasting/expansion operations. + # contiguous() is not enough: it keeps a zero stride on a size-1 dim. + if ( + isinstance(prop_constant_tensor, torch.Tensor) + and 0 in prop_constant_tensor.stride() + ): + prop_constant_tensor = prop_constant_tensor.clone( + memory_format=torch.contiguous_format + ) + # A folded constant must be a plain tensor that does not require # grad. no_grad is not enough when the op returns a view of a # parameter or the parameter itself. if isinstance(prop_constant_tensor, torch.Tensor): - prop_constant_tensor = without_zero_strides( - prop_constant_tensor - ).detach() + prop_constant_tensor = prop_constant_tensor.detach() const_node_to_tensor[node] = prop_constant_tensor # Only tensors can be lifted into placeholders. A non-tensor value read by @@ -458,11 +466,12 @@ def _update_output_node_and_specs( output_specs = exported_program.graph_signature.output_specs assert len(output_nodes) == len(output_specs) - clone_op = ( - exir_ops.edge.aten.clone.default - if _skip_dim_order - else exir_ops.edge.dim_order_ops._clone_dim_order.default - ) + if exported_program.dialect != "EDGE": + clone_op = torch.ops.aten.clone.default + elif _skip_dim_order: + clone_op = exir_ops.edge.aten.clone.default + else: + clone_op = exir_ops.edge.dim_order_ops._clone_dim_order.default for i in range(len(output_specs)): out_node = output_nodes[i] @@ -489,6 +498,8 @@ def constant_prop_pass( """ This pass is for constant propagation for Exported Program with lifted parameters, as the parameters will not be shown up as `get_attr` but as `placeholder` to the graph. + A graph output that is a constant gets a clone in front of it: an Edge clone when + exported_program.dialect is "EDGE", and torch.ops.aten.clone.default otherwise. Args: exported_program: The ExportedProgram to perform constant propagation on. diff --git a/exir/passes/propagate_device_pass.py b/exir/passes/propagate_device_pass.py index f0842ef788a..a3608b76de1 100644 --- a/exir/passes/propagate_device_pass.py +++ b/exir/passes/propagate_device_pass.py @@ -560,7 +560,8 @@ def _place_input_on_device( f"disable skip_h2d_for_method_inputs, or, for a mutable buffer, " f"have the partitioner put the buffer and every op that reads " f"or writes it in the delegate; tag_mutated_buffer tags the " - f"buffer once those ops are there." + f"buffer once those ops are there, unless the method also " + f"returns its new value or the new value is a method input." ) raise RuntimeError( f"{refusal} in place on the host. Either disable " diff --git a/exir/tests/test_delegate.py b/exir/tests/test_delegate.py index 915c832a591..4709c671762 100644 --- a/exir/tests/test_delegate.py +++ b/exir/tests/test_delegate.py @@ -274,6 +274,41 @@ def forward(self, x): torch.allclose(prog.exported_program().module()(*inputs), inputs[0] * 2) ) + def test_create_submodule_sorts_the_graph_in_place(self) -> None: + class Model(torch.nn.Module): + def forward(self, x, y): + a = x + y + b = a - y + c = a * y + return b / c + + inputs = (torch.randn(1, 3), torch.randn(1, 3)) + prog = to_edge(export(Model(), inputs, strict=True)) + gm = prog.exported_program().graph_module + graph = gm.graph + + node_list = [ + node + for node in graph.nodes + if node.op == "call_function" + and node.target + in {exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor} + ] + # The fused call goes after mul, so sub reads the add result before it + # is defined until the graph is sorted. + sub = next( + node for node in graph.nodes if node.target == exir_ops.edge.aten.sub.Tensor + ) + sub_gm, _ = create_submodule_from_nodes(gm, node_list, "tag") + + self.assertIs(gm.graph, graph) + self.assertIn(sub, graph.nodes) + graph.lint() + self.assertIsInstance(sub_gm(*inputs), list) + self.assertTrue( + torch.allclose(prog.exported_program().module()(*inputs), Model()(*inputs)) + ) + def test_create_submodule_list_return(self) -> None: """ Original graph: diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 81c5298aab1..ad626ec35e6 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -1637,6 +1637,35 @@ def forward(self) -> torch.Tensor: edge.exported_program()._validate() edge.to_executorch() + def test_constant_prop_for_output_clones_in_program_dialect(self) -> None: + class M(torch.nn.Module): + def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + return x * 2, torch.arange(4, dtype=torch.float32) + 1 + + exported = export(M(), (torch.ones(4),), strict=True) + for ep in (exported.run_decompositions({}), exported): + with self.subTest(dialect=ep.dialect): + ep = constant_prop_pass(ep) + [_, folded] = ep.graph.output_node().args[0] + self.assertEqual(folded.target, torch.ops.aten.clone.default) + ep._validate() + edge = to_edge( + ep, compile_config=EdgeCompileConfig(_skip_dim_order=False) + ) + [_, folded] = edge.exported_program().graph.output_node().args[0] + self.assertEqual( + folded.target, exir_ops.edge.dim_order_ops._clone_dim_order.default + ) + + with self.subTest(dialect="EDGE"): + edge = to_edge( + export(M(), (torch.ones(4),), strict=True), + compile_config=EdgeCompileConfig(_skip_dim_order=True), + ) + ep = constant_prop_pass(edge.exported_program()) + [_, folded] = ep.graph.output_node().args[0] + self.assertEqual(folded.target, exir_ops.edge.aten.clone.default) + def test_constant_prop_preserves_memory_alloc(self) -> None: class Add(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: