up

Author: metascroy

backends/mlx/CMakeLists.txt

@@ -212,25 +212,32 @@ set(MLX_METAL_JIT
     CACHE BOOL "Use JIT compiled Metal kernels"
 )
 
-# Auto-apply json patch so MLX reuses executorch's nlohmann_json instead of
-# downloading its own copy via FetchContent. TODO: upstream a patch to MLX
-set(MLX_JSON_PATCH "${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_json.patch")
-if(EXISTS "${MLX_JSON_PATCH}" AND EXISTS "${MLX_SOURCE_DIR}")
-  execute_process(
-    COMMAND git apply --check ${MLX_JSON_PATCH}
-    WORKING_DIRECTORY ${MLX_SOURCE_DIR}
-    RESULT_VARIABLE _mlx_json_patch_check
-    OUTPUT_QUIET ERROR_QUIET
-  )
-  if(_mlx_json_patch_check EQUAL 0)
+# Auto-apply patches to MLX submodule. Each patch is applied idempotently: `git
+# apply --check` tests whether the patch is still applicable (i.e. not yet
+# applied), and only then applies it.
+set(_mlx_patches
+    "${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_json.patch"
+    "${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_device_retain.patch"
+)
+foreach(_patch IN LISTS _mlx_patches)
+  if(EXISTS "${_patch}" AND EXISTS "${MLX_SOURCE_DIR}")
+    get_filename_component(_patch_name "${_patch}" NAME)
     execute_process(
-      COMMAND git apply ${MLX_JSON_PATCH} WORKING_DIRECTORY ${MLX_SOURCE_DIR}
+      COMMAND git apply --check "${_patch}"
+      WORKING_DIRECTORY ${MLX_SOURCE_DIR}
+      RESULT_VARIABLE _patch_check
+      OUTPUT_QUIET ERROR_QUIET
     )
-    message(STATUS "Applied mlx_json.patch to MLX submodule")
-  else()
-    message(STATUS "mlx_json.patch already applied or not applicable")
+    if(_patch_check EQUAL 0)
+      execute_process(
+        COMMAND git apply "${_patch}" WORKING_DIRECTORY ${MLX_SOURCE_DIR}
+      )
+      message(STATUS "Applied ${_patch_name} to MLX submodule")
+    else()
+      message(STATUS "${_patch_name} already applied or not applicable")
+    endif()
   endif()
-endif()
+endforeach()
 
 # Add MLX subdirectory
 message(STATUS "Adding MLX from submodule: ${MLX_SOURCE_DIR}")

backends/mlx/builder/program_builder.py

@@ -44,6 +44,7 @@
 )
 from executorch.backends.mlx.serialization.mlx_graph_schema import (
     FloatOrVid,
+    IdCopyNode,
     Instruction,
     InstructionChain,
     IntOrVid,
@@ -557,11 +558,51 @@ def check_support_only(self) -> None:
         # SymInts and corrupts the shape_env. This method is used for
         # ops_to_not_decompose() where we only need support status.
 
+    def _emit_buffer_mutation_writebacks(self):
+        """Emit copy-back instructions for BUFFER_MUTATION outputs.
+
+        When a model mutates a buffer (e.g., via .copy_() or .mul_()),
+        torch.export functionalizes it: the new value is a computation result,
+        and the output spec marks it as BUFFER_MUTATION with a target buffer.
+
+        This method emits an IdCopyNode for each BUFFER_MUTATION output,
+        copying the computation result back to the mutable buffer slot so
+        the updated value persists across execution calls.
+        """
+        from torch.export.graph_signature import InputKind, OutputKind
+
+        # Map buffer target name -> input placeholder name
+        target_to_placeholder = {}
+        for ispec in self.ep.graph_signature.input_specs:
+            if ispec.kind == InputKind.BUFFER and ispec.target is not None:
+                target_to_placeholder[ispec.target] = ispec.arg.name
+
+        for ospec in self.ep.graph_signature.output_specs:
+            if ospec.kind != OutputKind.BUFFER_MUTATION:
+                continue
+
+            result_slot = self.slot_manager.get_slot(ospec.arg.name)
+            placeholder_name = target_to_placeholder.get(ospec.target)
+            if result_slot is None or placeholder_name is None:
+                continue
+
+            buffer_slot = self.slot_manager.get_slot(placeholder_name)
+            if buffer_slot is None or buffer_slot.id_space != IdSpace.MutableBuffer:
+                continue
+
+            self.emit(
+                IdCopyNode(
+                    x=self.slot_to_tid(result_slot),
+                    out=self.slot_to_tid(buffer_slot),
+                )
+            )
+
     def build(self) -> MLXGraph:
         if self._mlx_graph is not None:
             return self._mlx_graph
 
         self._process_nodes()
+        self._emit_buffer_mutation_writebacks()
         self._verify_build()
         self._mlx_graph = self._build_mlx_graph()
         return self._mlx_graph

backends/mlx/patterns.py

@@ -36,7 +36,6 @@
 from executorch.backends.mlx.serialization.mlx_graph_schema import (
     AddIntNode,
     DequantizeNode,
-    IdCopyNode,
     IndexCopyNode,
     IntOrVid,
     IntOrVidOrTid,
@@ -125,19 +124,8 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot:
             )
         )
 
-        # index_copy returns the updated dst (same buffer, in-place update)
-        existing_slot = P.slot_manager.get_slot(n)
-        if existing_slot is not None and existing_slot != dst:
-            P.emit(
-                IdCopyNode(
-                    x=P.slot_to_tid(dst),
-                    out=P.slot_to_tid(existing_slot),
-                )
-            )
-            return existing_slot
-        else:
-            P.set_slot(n, dst)
-            return dst
+        P.set_slot(n, dst)
+        return dst
 
 
 @REGISTRY.register_pattern(name="ET_KV_CACHE_UPDATE")
@@ -238,21 +226,13 @@ def __call__(self, P: "MLXProgramBuilder", n: Node) -> Slot:
             [self.cache, self.update, self.start_pos]
         )
 
-        out_slot = P.make_or_get_slot(n)
-
         if self.ring_size > 0:
             self._emit_ring_buffer(P, cache_slot, update_slot, start_slot)
         else:
             self._emit_linear(P, cache_slot, update_slot, start_slot)
 
-        P.emit(
-            IdCopyNode(
-                x=P.slot_to_tid(cache_slot),
-                out=P.slot_to_tid(out_slot),
-            )
-        )
-
-        return out_slot
+        P.set_slot(n, cache_slot)
+        return cache_slot
 
     def _emit_linear(self, P: "MLXProgramBuilder", cache_slot, update_slot, start_slot):
         """Emit a single SliceUpdate for linear (non-ring) cache."""