diff --git a/backends/mlx/ops.py b/backends/mlx/ops.py index 5db12feef82..edd796e0b3f 100644 --- a/backends/mlx/ops.py +++ b/backends/mlx/ops.py @@ -94,6 +94,7 @@ GreaterEqualNode, GreaterNode, IdCopyNode, + IndexCopyNode, IntOrVid, IntOrVidOrTid, ItemIntNode, @@ -2406,6 +2407,38 @@ def _roll_handler(P: MLXProgramBuilder, n: Node) -> Slot: return out +@REGISTRY.register(target=[torch.ops.aten.index_copy.default]) +def _index_copy_handler(P: MLXProgramBuilder, n: Node) -> Slot: + """index_copy into a new tensor, leaving `dst` as it was. + + The INDEX_COPY pattern writes in place when `dst` is a buffer the method + mutates; anything else, a mutated user input included, lands here and the + caller writes the result back. + """ + args = P.args(n) + require_args(args, 4, 4, "aten.index_copy.default") + require_kwargs(P.kwargs(n), set(), "aten.index_copy.default") + dst, axis, indices, update = args + require_static_int(axis, "dim", "aten.index_copy.default") + # The runtime takes a 1-D index into a tensor with at least one dim; ATen also + # allows 0-d ones, which stay off the delegate. + if n.args[0].meta["val"].dim() == 0 or n.args[2].meta["val"].dim() != 1: + raise NotImplementedError( + "aten.index_copy.default needs a 1-D index into a tensor of rank >= 1" + ) + out = P.make_or_get_slot(n) + P.emit( + IndexCopyNode( + dst=P.slot_to_tid(dst), + update=P.slot_to_tid(update), + indices=P.slot_to_tid(indices), + out=P.slot_to_tid(out), + axis=axis, + ) + ) + return out + + @REGISTRY.register(target=[torch.ops.aten.index.Tensor]) def _index_handler(P: MLXProgramBuilder, n: Node) -> Slot: args = P.args(n) diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py index 8abd8933e4f..c1d682f9a1c 100644 --- a/backends/mlx/patterns.py +++ b/backends/mlx/patterns.py @@ -123,6 +123,10 @@ def _int4_qparams_lower(qdata_node: Node, scale_node: Node, group_size: int) -> class IndexCopyHandler(PatternHandler): """ Pattern for index-based updates on mutable buffers. + + Only a buffer the method mutates is written in place: the delegate owns it + across calls. A mutated user input is not the delegate's to keep, so its + index_copy is left to the functional handler and written back by the caller. """ def __init__( @@ -148,10 +152,8 @@ def maybe_create( # noqa: C901 if not match_target(index_copy_node, torch.ops.aten.index_copy.default): return None - # index_copy should write to a mutable input/buffer to be an index update. - if (index_copy_node.name not in ep.graph_signature.buffers_to_mutate) and ( - index_copy_node.name not in ep.graph_signature.user_inputs_to_mutate - ): + # index_copy should write to a mutable buffer to be an index update. + if index_copy_node.name not in ep.graph_signature.buffers_to_mutate: return None # index_copy(dst, axis, indices, update) diff --git a/backends/mlx/preprocess.py b/backends/mlx/preprocess.py index b2e26a4d2b2..81e89743018 100644 --- a/backends/mlx/preprocess.py +++ b/backends/mlx/preprocess.py @@ -57,6 +57,11 @@ class MLXBackend(BackendDetails): EXPECTED_MAGIC: ClassVar[bytes] = MAGIC EXPECTED_LENGTH: ClassVar[int] = HEADER_LENGTH + # MLXBackend::execute releases an input that is also an output before + # evaluating, so MLX can update it in its own memory, and only writes + # outputs once everything reading inputs is done. + writes_mutated_inputs_in_place: ClassVar[bool] = True + @staticmethod def preprocess( edge_program: ExportedProgram, diff --git a/backends/mlx/runtime/MLXBackend.cpp b/backends/mlx/runtime/MLXBackend.cpp index a30abf2a185..e64a717c867 100644 --- a/backends/mlx/runtime/MLXBackend.cpp +++ b/backends/mlx/runtime/MLXBackend.cpp @@ -32,10 +32,13 @@ #include +#include +#include #include #include #include #include +#include #include @@ -128,9 +131,21 @@ array prepare_output( return result; } -// Wait for a prepared output array and copy its data to an ET tensor. +// Whether [a, a + a_nbytes) and [b, b + b_nbytes) share any byte. +bool overlaps(const void* a, size_t a_nbytes, const void* b, size_t b_nbytes) { + const auto a_begin = reinterpret_cast(a); + const auto b_begin = reinterpret_cast(b); + return a_nbytes > 0 && b_nbytes > 0 && a_begin < b_begin + b_nbytes && + b_begin < a_begin + a_nbytes; +} + +// Wait for a prepared output array and copy its data to an ET tensor, or copy +// `snapshot` instead when it holds the array's data taken earlier. // The array must have been submitted via async_eval before calling this. -void write_output(array& arr, ETTensor& out) { +void write_output( + array& arr, + ETTensor& out, + const std::vector& snapshot = {}) { arr.wait(); // Resize output tensor if shape doesn't match (dynamic shapes) @@ -171,12 +186,16 @@ void write_output(array& arr, ETTensor& out) { " bytes, output has " + std::to_string(out_nbytes) + " bytes"); } - const void* src = arr.data(); + const void* src = snapshot.empty() ? arr.data() : snapshot.data(); if (!src) { throw std::runtime_error( "write_output: arr.data() is null after wait()"); } - std::memcpy(out.mutable_data_ptr(), src, out_nbytes); + // An output that is also an input can be computed into the input's own + // memory when MLX donates its buffer, and is then already in place. + if (src != out.const_data_ptr()) { + std::memcpy(out.mutable_data_ptr(), src, out_nbytes); + } } } // namespace @@ -472,6 +491,7 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { }; std::vector tensor_output_info; + std::vector inputs_written_in_place; size_t arg_idx = 0; auto* h = static_cast(handle); @@ -504,6 +524,13 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { ")"); } if (slot.slot_type == SlotType::TensorSlot) { + // The method passes an input it mutates in place as the output + // that carries the mutation too, as one and the same EValue. + EValue* arg = args[arg_idx]; + if (std::find(args.begin() + n_inputs, args.end(), arg) != + args.end()) { + inputs_written_in_place.push_back(Tid{slot.idx}); + } const ETTensor& tensor = args[arg_idx++]->toTensor(); Tid tid{slot.idx}; std::optional expected_meta = std::nullopt; @@ -572,6 +599,12 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { } } + // With nothing but the graph left holding an input that is written in + // place, MLX can donate the input's buffer to its update. + for (Tid tid : inputs_written_in_place) { + h->state.release_tensor(tid); + } + // Submit all output work to GPU asynchronously // async_eval encodes Metal commands and returns immediately. // The GPU will signal events on completion. @@ -581,11 +614,50 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { } // Lock released — GPU is still executing - for (auto& info : tensor_output_info) { + // Per tensor output, a copy of its data when it must not be read from + // the array at the time it is written. + std::vector> snapshots(tensor_output_info.size()); + if (!inputs_written_in_place.empty()) { + // An output written into an input's memory must not land while work + // still reading that input is in flight. + for (auto& arr : prepared_outputs) { + arr.wait(); + } + // Outputs are then written one after another, so one written into an + // input's memory would also change any output still reading that + // memory and written after it, e.g. a clone of the input taken before + // the update, which MLX leaves as the input's own buffer. Copy such + // outputs out before writing any. + for (size_t i = 0; i < tensor_output_info.size(); ++i) { + const array& arr = + prepared_outputs[tensor_output_info[i].prepared_idx]; + // Outputs written up to and including this one. + for (size_t w = 0; w <= i; ++w) { + const auto& written = tensor_output_info[w]; + const array& written_arr = prepared_outputs[written.prepared_idx]; + const void* dst = + args[written.arg_idx]->toTensor().const_data_ptr(); + if (written_arr.data() != dst && + overlaps( + arr.data(), + arr.nbytes(), + dst, + written_arr.nbytes())) { + const auto* src = static_cast(arr.data()); + snapshots[i].assign(src, src + arr.nbytes()); + break; + } + } + } + } + + for (size_t i = 0; i < tensor_output_info.size(); ++i) { + const auto& info = tensor_output_info[i]; ETTensor& out_tensor = args[info.arg_idx]->toTensor(); // write_output waits on arr to be ready - write_output(prepared_outputs[info.prepared_idx], out_tensor); + write_output( + prepared_outputs[info.prepared_idx], out_tensor, snapshots[i]); } h->state.reset(); // Release temp GPU buffers back to MLX cache diff --git a/backends/mlx/runtime/MLXExecutor.h b/backends/mlx/runtime/MLXExecutor.h index 95414ae8e5c..46f85acdb22 100644 --- a/backends/mlx/runtime/MLXExecutor.h +++ b/backends/mlx/runtime/MLXExecutor.h @@ -320,6 +320,10 @@ struct ExecutionState { return idx; } + void release_tensor(Tid id) { + tensors[tensor_index(id)] = std::nullopt; + } + void reset() { // Clear per-execution tensors (inputs, outputs, temps) // Constants and mutable buffers are not in this vector diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index 135baa2dcad..3614e0fb83e 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -1317,26 +1317,26 @@ exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) { const size_t uaxis = static_cast(axis); const int dst_dim = static_cast(dst.shape()[uaxis]); - // Get indices as a vector of ints, handling negative indices - // Note: PyTorch uses int64 for indices, so we read as int64_t - eval(indices); // Ensure indices are materialized before accessing data - if (indices.dtype() != ::mlx::core::int64) { + // Like ATen, negative indices are rejected rather than wrapped. Export + // accepts int32 indices too, so they are widened rather than rejected. + const array idx_arr = indices.dtype() == ::mlx::core::int32 + ? astype(indices, ::mlx::core::int64, s) + : indices; + eval(idx_arr); // Ensure indices are materialized before accessing data + if (idx_arr.dtype() != ::mlx::core::int64) { throw std::invalid_argument( - std::string("IndexCopyNode: expected int64 indices, got ") + + std::string("IndexCopyNode: expected int64 or int32 indices, got ") + ExecutionState::dtype_str(indices.dtype())); } - std::vector idx_vec(indices.size()); - auto idx_data = indices.data(); - for (size_t i = 0; i < indices.size(); ++i) { - int64_t idx = idx_data[i]; - if (idx < 0) { - idx += dst_dim; - } + std::vector idx_vec(idx_arr.size()); + auto idx_data = idx_arr.data(); + for (size_t i = 0; i < idx_arr.size(); ++i) { + const int64_t idx = idx_data[i]; if (idx < 0 || idx >= dst_dim) { throw std::out_of_range( - "IndexCopyNode: index " + std::to_string(idx_data[i]) + - " out of range for axis " + std::to_string(axis) + " with size " + - std::to_string(dst_dim)); + "index_copy: index " + std::to_string(idx) + + " is out of bounds for dimension " + std::to_string(axis) + + " with size " + std::to_string(dst_dim)); } if (idx > std::numeric_limits::max()) { throw std::out_of_range( diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index d8b7b8e485a..9649769de87 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -19,7 +19,13 @@ from executorch.backends.mlx.partitioner import MLXPartitioner 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 import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge, + to_edge_transform_and_lower, +) +from executorch.exir.passes import MemoryPlanningPass from executorch.runtime import Runtime from torch.export import export @@ -428,5 +434,251 @@ def forward(self, x): self.assertLess(err, 1e-4) +class TestMLXPartitionerMutatedUserInput(unittest.TestCase): + """index_copy_ into a user input lowers functionally and is written back.""" + + class IndexCopyInput(nn.Module): + def __init__(self, axis): + super().__init__() + self.axis = axis + + def forward(self, data, indices, update): + data.index_copy_(self.axis, indices, update) + return data.sum(self.axis) + + def test_index_copy_into_input_lowers_and_executes(self): + # The write-back of the mutated input runs outside the delegate, as a + # portable copy_, so this goes through the Python runtime rather than + # op_test_runner, which links no portable kernels. + for axis, shape, update_shape in ( + (0, (16, 8), (4, 8)), + (1, (8, 16), (8, 4)), + ): + with self.subTest(axis=axis): + model = self.IndexCopyInput(axis).eval() + inputs = ( + torch.randn(shape), + torch.tensor([1, 5, 9, 12]), + torch.randn(update_shape), + ) + expected = inputs[0].clone() + with torch.no_grad(): + expected_sum = model(expected, *inputs[1:]) + program = _lower(model, inputs) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + counts = get_mlx_node_counts(path) + method = Runtime.get().load_program(path).load_method("forward") + data, summed = method.execute(list(inputs)) + self.assertEqual(_delegate_count(program), 1) + self.assertEqual(counts.get("IndexCopyNode", 0), 1) + self.assertTrue(torch.allclose(data, expected)) + self.assertTrue(torch.allclose(summed, expected_sum, atol=1e-5)) + + def test_index_copy_into_input_written_in_place_with_reinplace(self): + # With run_reinplace_pass the delegate writes the mutated input itself: + # the method has no write-back copy_, and returns the input value. + model = self.IndexCopyInput(0).eval() + cache = torch.zeros(16, 8) + program = to_edge_transform_and_lower( + export(model, (cache, torch.tensor([1]), torch.randn(1, 8)), strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + self.assertFalse( + any(n.target == torch.ops.aten.copy_.default for n in graph.nodes) + ) + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.outputs[0], plan.inputs[0]) + + method = Runtime.get().load_program(program.buffer).load_method("forward") + expected = cache.clone() + for row in (3, 7, 11): + indices, update = torch.tensor([row]), torch.randn(1, 8) + expected.index_copy_(0, indices, update) + cache, summed = method.execute([cache, indices, update]) + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.allclose(summed, expected.sum(0), atol=1e-5)) + + class ReadThenIndexCopy(nn.Module): + def forward(self, data, indices, update): + before = data.sum(0) + data.index_copy_(0, indices, update) + return before, data.sum(0) + + def test_read_before_index_copy_with_reinplace(self): + # A value read from the input before it is updated in place must be + # computed from the old contents, not race the write into the input. + model = self.ReadThenIndexCopy().eval() + cache = torch.randn(16, 8) + program = to_edge_transform_and_lower( + export(model, (cache, torch.tensor([1]), torch.randn(1, 8)), strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + self.assertFalse( + any(n.target == torch.ops.aten.copy_.default for n in graph.nodes) + ) + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.outputs[0], plan.inputs[0]) + + method = Runtime.get().load_program(program.buffer).load_method("forward") + expected = cache.clone() + for row in (3, 7, 11): + indices, update = torch.tensor([row]), torch.randn(1, 8) + expected_before = expected.sum(0) + expected.index_copy_(0, indices, update) + cache, before, after = method.execute([cache, indices, update]) + self.assertTrue(torch.allclose(before, expected_before, atol=1e-5)) + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.allclose(after, expected.sum(0), atol=1e-5)) + + def test_unplanned_input_written_in_place_with_reinplace(self): + # An unplanned input is passed to MLX without a copy when its memory can + # be wrapped, so the update lands in the caller's own tensor. A large + # cache is page-aligned, which wrapping without a copy may require. + model = self.IndexCopyInput(0).eval() + cache = torch.zeros(1024, 1024) + program = to_edge_transform_and_lower( + export( + model, + (cache.clone(), torch.tensor([1]), torch.randn(1, 1024)), + strict=False, + ), + partitioner=[MLXPartitioner()], + ).to_executorch( + ExecutorchBackendConfig( + run_reinplace_pass=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ) + ) + method = Runtime.get().load_program(program.buffer).load_method("forward") + expected = cache.clone() + for row in (3, 700): + indices, update = torch.tensor([row]), torch.randn(1, 1024) + expected.index_copy_(0, indices, update) + _, summed = method.execute([cache, indices, update]) + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.allclose(summed, expected.sum(0), atol=1e-3)) + + class TwoCaches(nn.Module): + def forward(self, k_cache, v_cache, indices, k, v): + k_cache.index_copy_(0, indices, k) + v_cache.index_copy_(0, indices, v) + return k_cache.sum(0) + v_cache.sum(0) + + def test_two_inputs_written_in_place_with_reinplace(self): + # Two caches mutated by one delegate are both written in place. + model = self.TwoCaches().eval() + k_cache, v_cache = torch.zeros(16, 8), torch.zeros(16, 8) + example = ( + k_cache, + v_cache, + torch.tensor([1]), + torch.randn(1, 8), + torch.randn(1, 8), + ) + program = to_edge_transform_and_lower( + export(model, tuple(t.clone() for t in example), strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + self.assertFalse( + any(n.target == torch.ops.aten.copy_.default for n in graph.nodes) + ) + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.outputs[:2], plan.inputs[:2]) + + method = Runtime.get().load_program(program.buffer).load_method("forward") + expected_k, expected_v = k_cache.clone(), v_cache.clone() + for row in (3, 7): + indices, k, v = torch.tensor([row]), torch.randn(1, 8), torch.randn(1, 8) + expected_k.index_copy_(0, indices, k) + expected_v.index_copy_(0, indices, v) + k_cache, v_cache, summed = method.execute([k_cache, v_cache, indices, k, v]) + self.assertTrue(torch.equal(k_cache, expected_k)) + self.assertTrue(torch.equal(v_cache, expected_v)) + expected = expected_k.sum(0) + expected_v.sum(0) + self.assertTrue(torch.allclose(summed, expected, atol=1e-5)) + + class SnapshotBeforeMutation(nn.Module): + def forward(self, cache, indices, update): + before = cache.clone() + cache.index_copy_(0, indices, update) + return before, cache + + class ViewOfSnapshotBeforeMutation(nn.Module): + def forward(self, cache, indices, update): + before = cache.clone() + cache.index_copy_(0, indices, update) + return before.view(2, 2), cache + + def test_snapshot_before_index_copy_with_reinplace(self): + # MLX can return the clone as the input's own memory, so it must be + # read before the update is written into that memory. The view puts + # the snapshot after the update in the delegate's outputs, so the + # update is written first. + for model in ( + self.SnapshotBeforeMutation().eval(), + self.ViewOfSnapshotBeforeMutation().eval(), + ): + with self.subTest(model=type(model).__name__): + cache = torch.tensor([1.0, 2.0, 3.0, 4.0]) + indices = torch.tensor([1, 3], dtype=torch.int64) + update = torch.tensor([20.0, 40.0]) + program = to_edge_transform_and_lower( + export(model, (cache.clone(), indices, update), strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + self.assertFalse( + any(n.target == torch.ops.aten.copy_.default for n in graph.nodes) + ) + plan = program.executorch_program.execution_plan[0] + self.assertEqual(plan.outputs[0], plan.inputs[0]) + + method = ( + Runtime.get().load_program(program.buffer).load_method("forward") + ) + expected = cache.clone() + for step in range(3): + update = torch.tensor([20.0, 40.0]) * (step + 1) + expected_before = expected.clone() + expected.index_copy_(0, indices, update) + cache, before, after = method.execute([cache, indices, update]) + self.assertTrue(torch.equal(before.flatten(), expected_before)) + self.assertTrue(torch.equal(cache, expected)) + self.assertTrue(torch.equal(after, expected)) + + def test_index_copy_accepts_int32_indices(self): + # Export accepts int32 indices for index_copy, so the delegate must too. + # Eager ATen does not, so the reference uses int64 indices. + model = self.IndexCopyInput(0).eval() + data, update = torch.randn(16, 8), torch.randn(2, 8) + indices = torch.tensor([2, 9], dtype=torch.int32) + expected = data.clone().index_copy_(0, indices.long(), update) + program = _lower(model, (data.clone(), indices, update)) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + method = Runtime.get().load_program(path).load_method("forward") + result, summed = method.execute([data.clone(), indices, update]) + self.assertTrue(torch.equal(result, expected)) + self.assertTrue(torch.allclose(summed, expected.sum(0), atol=1e-5)) + + def test_index_copy_rejects_negative_index(self): + # ATen raises on negative indices, so the delegate must not wrap them. + model = self.IndexCopyInput(0).eval() + inputs = (torch.randn(4, 8), torch.tensor([1]), torch.randn(1, 8)) + program = _lower(model, inputs) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + method = Runtime.get().load_program(path).load_method("forward") + with self.assertRaises(RuntimeError): + method.execute([inputs[0], torch.tensor([-1]), inputs[2]]) + + if __name__ == "__main__": unittest.main() diff --git a/exir/backend/backend_details.py b/exir/backend/backend_details.py index ca7d3b6e66a..025f288023d 100644 --- a/exir/backend/backend_details.py +++ b/exir/backend/backend_details.py @@ -104,6 +104,16 @@ class BackendDetails(ABC): """ + # Whether the runtime backend computes its outputs correctly when the method + # passes it the same EValue as one of its inputs and one of its outputs. + # With run_reinplace_pass, an input the delegate mutates is then written in + # place, instead of copied back after the delegate returns. Opting in means + # the backend reads that input before writing the aliased output, and keeps + # the contents of any other output whose storage overlaps it. EXIR still only + # aliases a mutated placeholder that nothing reads between the delegate and + # its write-back, and that the delegate doesn't also take through an alias. + writes_mutated_inputs_in_place: bool = False + def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) diff --git a/exir/capture/_config.py b/exir/capture/_config.py index 13ab5b3f46c..7cfa6dfefca 100644 --- a/exir/capture/_config.py +++ b/exir/capture/_config.py @@ -121,7 +121,9 @@ class ExecutorchBackendConfig: # If set to true, we run quant fusion and constant propagation passes do_quant_fusion_and_const_prop: bool = False - # Experimental: If set to true, we run a pass to reinplace ops in the graph. + # Experimental: If set to true, we run a pass to reinplace ops in the graph, + # and let a delegate whose backend sets writes_mutated_inputs_in_place + # write a mutated input in place instead of copying it back. run_reinplace_pass: bool = False # When True, memory planning partitions specs by device and runs the diff --git a/exir/passes/reinplace.py b/exir/passes/reinplace.py index b5da4cc605c..51465ed8477 100644 --- a/exir/passes/reinplace.py +++ b/exir/passes/reinplace.py @@ -6,6 +6,7 @@ # pyre-strict +import operator from typing import Any, Dict, FrozenSet, Iterable, Optional, Set, Tuple, Union import torch @@ -15,6 +16,7 @@ unwrap_op_overload, ) from torch.export import ExportedProgram +from torch.export.graph_signature import ExportGraphSignature # --------------------------------------------------------------------------- @@ -530,3 +532,82 @@ def reinplace_pass( # noqa: C901 spec.arg.name = arg.name return ep + + +def _storage_aliases(node: torch.fx.Node) -> Set[torch.fx.Node]: + """``node`` and every node sharing its storage, transitively: views, which + ReplaceViewCopyWithViewPass has turned into ``memory.view``, and results + planned onto the TensorSpec of the alias they consume, like those of + in-place ops, including on views.""" + from executorch.exir import memory + + aliases = {node} + frontier = [node] + while frontier: + alias = frontier.pop() + spec = alias.meta.get("spec") + for user in alias.users: + if user in aliases: + continue + if (user.target == memory.view and user.args[0] is alias) or ( + spec is not None and user.meta.get("spec") is spec + ): + aliases.add(user) + frontier.append(user) + return aliases + + +def reinplace_delegate_input_mutations( + graph_module: torch.fx.GraphModule, graph_signature: ExportGraphSignature +) -> None: + """Let a delegate write an input it mutates in place. + + The method writes a mutated input back with ``copy_(input, output)`` once + the delegate that computed ``output`` returns. When that delegate also + reads ``input`` and its backend sets ``writes_mutated_inputs_in_place``, + give ``output`` the input's TensorSpec, so the two are one value the + delegate writes directly, and drop the copy. + + Runs on a graph with tensor specs, before memory planning. + """ + from executorch.exir.backend.backend_details import BackendDetails + from executorch.exir.delegate import executorch_call_delegate + + backends = {cls.__name__: cls for cls in BackendDetails.__subclasses__()} + graph = graph_module.graph + position = {node: i for i, node in enumerate(graph.nodes)} + for node in list(graph.nodes): + if node.target != torch.ops.aten.copy_.default: + continue + mutated, update = node.args[0], node.args[1] + if mutated.op != "placeholder" or update.target != operator.getitem: + continue + delegate = update.args[0] + if delegate.target != executorch_call_delegate or mutated not in delegate.args: + continue + backend = backends.get( + getattr(graph_module, delegate.args[0].target).backend_id + ) + if backend is None or not backend.writes_mutated_inputs_in_place: + continue + # The input now changes while the delegate runs, not at the copy, so + # nothing in between may read it, directly or through an alias, and the + # delegate may not also take it through an alias. + aliases = _storage_aliases(mutated) + if any( + position[delegate] < position[user] < position[node] + for alias in aliases + for user in alias.users + ) or any(alias is not mutated and alias in delegate.args for alias in aliases): + continue + # The emitter takes the delegate's outputs from its own specs. + delegate_specs = list(delegate.meta["spec"]) + delegate_specs[update.args[1]] = mutated.meta["spec"] + delegate.meta["spec"] = type(delegate.meta["spec"])(delegate_specs) + update.meta["spec"] = mutated.meta["spec"] + for spec in graph_signature.output_specs: + if spec.arg.name == node.name: + spec.arg.name = update.name + node.replace_all_uses_with(update) + graph.erase_node(node) + graph_module.recompile() diff --git a/exir/program/_program.py b/exir/program/_program.py index b1d896281c0..09705b8d3a7 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -67,7 +67,11 @@ from executorch.exir.passes.propagate_device_config import PropagateDeviceConfig from executorch.exir.passes.propagate_device_pass import PropagateDevicePass from executorch.exir.passes.quant_fusion_pass import quant_fusion_and_const_prop_pass -from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS, reinplace_pass +from executorch.exir.passes.reinplace import ( + DEFAULT_INPLACEABLE_OPS, + reinplace_delegate_input_mutations, + reinplace_pass, +) from executorch.exir.passes.remove_graph_asserts_pass import ( RemoveGraphAssertsPass, RemoveNonCoreAtenOpGraphAssertsPass, @@ -1770,6 +1774,9 @@ def to_executorch( # noqa (FLAKE8) C901 # TODO(who?) p.update_placeholder_tensor_specs(program, new_gm) + if config.run_reinplace_pass: + reinplace_delegate_input_mutations(new_gm, new_signature) + # Tag constant weights. if ( isinstance(config.external_constants, bool) diff --git a/exir/tests/test_reinplace_pass.py b/exir/tests/test_reinplace_pass.py index 52675fedb59..a214e694475 100644 --- a/exir/tests/test_reinplace_pass.py +++ b/exir/tests/test_reinplace_pass.py @@ -10,10 +10,24 @@ from typing import List, Optional import torch -from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir import ( + EdgeCompileConfig, + memory, + to_edge, + to_edge_transform_and_lower, +) +from executorch.exir.backend.backend_details import BackendDetails, PreprocessResult +from executorch.exir.backend.canonical_partitioners.all_node_partitioner import ( + AllNodePartitioner, +) from executorch.exir.capture._config import ExecutorchBackendConfig +from executorch.exir.delegate import executorch_call_delegate from executorch.exir.dialects._ops import ops as edge_ops -from executorch.exir.passes.reinplace import DEFAULT_INPLACEABLE_OPS, reinplace_pass +from executorch.exir.passes.reinplace import ( + DEFAULT_INPLACEABLE_OPS, + reinplace_delegate_input_mutations, + reinplace_pass, +) from executorch.extension.pybindings.portable_lib import ( # @manual=//executorch/extension/pybindings:portable_lib _load_for_executorch_from_buffer, ) @@ -40,6 +54,28 @@ def _find_nodes( ] +class _InPlaceDemoBackend(BackendDetails): + writes_mutated_inputs_in_place = True + + @staticmethod + def preprocess(edge_program, compile_specs) -> PreprocessResult: + return PreprocessResult(processed_bytes=b"") + + +class _CopyingDemoBackend(BackendDetails): + @staticmethod + def preprocess(edge_program, compile_specs) -> PreprocessResult: + return PreprocessResult(processed_bytes=b"") + + +class _CacheUpdate(torch.nn.Module): + def forward( + self, k: torch.Tensor, k_cache: torch.Tensor, pos: torch.Tensor + ) -> torch.Tensor: + k_cache.index_copy_(2, pos, k) + return k_cache.sum(-1) + + class TestReinplacePass(unittest.TestCase): def test_index_put_reinplace(self) -> None: """Test that index_put on a mutable buffer can be reinplaced.""" @@ -813,6 +849,76 @@ def forward( # program). edge.to_executorch() + def test_delegate_writes_mutated_input_in_place(self) -> None: + """A delegate that mutates one of its inputs writes it in place when + its backend opts in: its output shares the input's spec and the + write-back copy_ is gone. Other backends keep the copy_.""" + inputs = (torch.ones(1, 2, 1, 4), torch.zeros(1, 2, 8, 4), torch.tensor([3])) + for backend, in_place in ( + ("_InPlaceDemoBackend", True), + ("_CopyingDemoBackend", False), + ): + with self.subTest(backend=backend): + program = to_edge_transform_and_lower( + export(_CacheUpdate(), inputs, strict=True), + partitioner=[AllNodePartitioner(backend, [])], + ).to_executorch(ExecutorchBackendConfig(run_reinplace_pass=True)) + graph = program.exported_program().graph + k_cache = next(n for n in graph.nodes if n.name == "k_cache") + delegate = next( + n for n in graph.nodes if n.target == executorch_call_delegate + ) + self.assertEqual( + len(_find_nodes(program.exported_program(), "copy_")), + 0 if in_place else 1, + ) + self.assertEqual( + any(spec is k_cache.meta["spec"] for spec in delegate.meta["spec"]), + in_place, + ) + + def test_delegate_input_mutation_kept_when_read_before_write_back( + self, + ) -> None: + """Writing the input in place moves its mutation from the copy_ to the + delegate, so anything reading the input in between, directly or through + a view taken before the delegate, would see the new values. The copy_ + then stays and the specs stay separate.""" + inputs = (torch.ones(1, 2, 1, 4), torch.zeros(1, 2, 8, 4), torch.tensor([3])) + for read_through in ("input", "view", "in-place op on a view"): + with self.subTest(read_through=read_through): + program = to_edge_transform_and_lower( + export(_CacheUpdate(), inputs, strict=True), + partitioner=[AllNodePartitioner("_InPlaceDemoBackend", [])], + ).to_executorch() + ep = program.exported_program() + graph = ep.graph + k_cache = next(n for n in graph.nodes if n.name == "k_cache") + delegate = next( + n for n in graph.nodes if n.target == executorch_call_delegate + ) + read = k_cache + with graph.inserting_before(delegate): + if read_through != "input": + read = graph.call_function(memory.view, (k_cache, [16, 4])) + read.meta["spec"] = object() + if read_through == "in-place op on a view": + # Planned onto the view's spec, not the input's. + update_view = graph.call_function( + torch.ops.aten.add_.Tensor, (read, read) + ) + update_view.meta["spec"] = read.meta["spec"] + read = update_view + with graph.inserting_after(delegate): + graph.call_function(torch.ops.aten.clone.default, (read,)) + + reinplace_delegate_input_mutations(ep.graph_module, ep.graph_signature) + + self.assertEqual(len(_find_nodes(ep, "copy_")), 1) + self.assertFalse( + any(spec is k_cache.meta["spec"] for spec in delegate.meta["spec"]) + ) + def test_chain_of_inplaceable_ops(self) -> None: """A chain of safe-to-reinplace ops gets fully rewritten in topological order. Exercises: