MLX: lower index_copy into a mutated user input functionally

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)

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)

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,

backends/mlx/runtime/MLXBackend.cpp

@@ -32,10 +32,13 @@
 
 #include <executorch/backends/mlx/runtime/backend_options.h>
 
+#include <algorithm>
+#include <cstdint>
 #include <cstring>
 #include <limits>
 #include <memory>
 #include <mutex>
+#include <vector>
 
 #include <TargetConditionals.h>
 
@@ -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<uintptr_t>(a);
+  const auto b_begin = reinterpret_cast<uintptr_t>(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<uint8_t>& 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<void>();
+  const void* src = snapshot.empty() ? arr.data<void>() : snapshot.data();
   if (!src) {
     throw std::runtime_error(
         "write_output: arr.data<void>() 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<OutputInfo> tensor_output_info;
+      std::vector<Tid> inputs_written_in_place;
       size_t arg_idx = 0;
 
       auto* h = static_cast<MLXHandle*>(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<TensorMeta> 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<std::vector<uint8_t>> 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<void>() != dst &&
+                overlaps(
+                    arr.data<void>(),
+                    arr.nbytes(),
+                    dst,
+                    written_arr.nbytes())) {
+              const auto* src = static_cast<const uint8_t*>(arr.data<void>());
+              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

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

backends/mlx/runtime/MLXInterpreter.h

@@ -1308,7 +1308,11 @@ inline std::tuple<int, int, int, int> next_contiguous_run(
 // Optimizes into slice_update calls for contiguous runs
 inline void
 exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) {
-  array& dst = st.tensor_ref(n.dst);
+  // When out == dst, use direct assignment to preserve MLX buffer donation.
+  // Only then is dst written, so a functional update may read a constant.
+  const bool in_place = (n.out.idx == n.dst.idx);
+  const array& dst =
+      in_place ? st.tensor_ref(n.dst) : st.const_tensor_ref(n.dst);
   const array& upd = st.const_tensor_ref(n.update);
   const array& indices = st.const_tensor_ref(n.indices);
   if (indices.ndim() != 1) {
@@ -1319,26 +1323,31 @@ exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) {
   const size_t uaxis = static_cast<size_t>(axis);
   const int dst_dim = static_cast<int>(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.
+  // Indices computed in the graph can be a strided view, so they are made
+  // contiguous before their data is read.
+  const array idx_arr = contiguous(
+      indices.dtype() == ::mlx::core::int32
+          ? astype(indices, ::mlx::core::int64, s)
+          : indices,
+      false,
+      s);
+  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<int32_t> idx_vec(indices.size());
-  auto idx_data = indices.data<int64_t>();
-  for (size_t i = 0; i < indices.size(); ++i) {
-    int64_t idx = idx_data[i];
-    if (idx < 0) {
-      idx += dst_dim;
-    }
+  std::vector<int32_t> idx_vec(idx_arr.size());
+  auto idx_data = idx_arr.data<int64_t>();
+  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<int32_t>::max()) {
       throw std::out_of_range(
@@ -1348,11 +1357,9 @@ exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) {
     idx_vec[i] = static_cast<int>(idx);
   }
 
-  // When out == dst, use direct assignment to preserve MLX buffer donation.
-  // TODO: I'm not sure if this is needed as a special case since the standard
-  // st.set_tensor does a std::move. Keeping for now, but should investigate and
-  // possibly remove in future.
-  const bool in_place = (n.out.idx == n.dst.idx);
+  // TODO: I'm not sure the in-place case is needed as a special case since the
+  // standard st.set_tensor does a std::move. Keeping for now, but should
+  // investigate and possibly remove in future.
 
   if (idx_vec.empty()) {
     if (!in_place) {
@@ -1402,8 +1409,9 @@ exec_index_copy(const IndexCopyNode& n, ExecutionState& st, StreamOrDevice s) {
         : slice(upd, to_shape(upd_vstart), to_shape(upd_vstop), s);
 
     if (in_place) {
-      dst = slice_update(
-          dst, upd_slice, to_shape(dst_vstart), to_shape(dst_vstop), s);
+      array& target = st.tensor_ref(n.dst);
+      target = slice_update(
+          target, upd_slice, to_shape(dst_vstart), to_shape(dst_vstop), s);
     } else {
       result = slice_update(
           result, upd_slice, to_shape(dst_vstart), to_shape(dst_vstop), s);