Sync upstream ml-explore/mlx-swift 392aa34 (mlx v0.32.4)

.github/scripts/build-linux-cuda-cmake.sh

@@ -17,13 +17,6 @@ mkdir -p build
 pushd build
 cmake -DMLX_BUILD_METAL=OFF -DMLX_BUILD_CUDA=ON -DMLX_C_BUILD_EXAMPLES=OFF .. -G Ninja
 
-# Verify the local mlx patch is actually in place: without it the examples below
-# crash on exit, which is easy to misread as an unrelated CUDA failure.
-git -C _deps/mlx-src apply --reverse --check "$PWD/../cmake/mlx.patch" || {
-  echo "error: cmake/mlx.patch is not applied to _deps/mlx-src"
-  exit 1
-}
-
 ninja
 ./example1 --device gpu
 ./tutorial --device gpu

CMakeLists.txt

@@ -37,22 +37,10 @@ endif()
 set(apply_patch ${CMAKE_COMMAND} -DREPO=<SOURCE_DIR>)
 set(apply_patch_script -P ${CMAKE_CURRENT_SOURCE_DIR}/cmake/apply-patch.cmake)
 
-# two local patches for CUDA:
-#
-# * https://github.com/ml-explore/mlx/pull/4480: leak the global command encoder
-#   map rather than synchronizing on process shutdown
-# * make ~CudaHandle non-throwing: cu::Worker is detached and holds the last
-#   reference to itself, so ~Worker -> ~CudaEvent -> cudaEventDestroy runs on
-#   the worker thread at an arbitrary point during exit, racing the CUDA
-#   runtime's own teardown. A failing destroy there terminated the process.
-set(mlx_patch
-    ${apply_patch} -DPATCH=${CMAKE_CURRENT_SOURCE_DIR}/cmake/mlx.patch
-    ${apply_patch_script})
-
 # note: this must be declared (and named `mlx`) before mlx-c is made available
-# so that our pinned/patched version wins over the one mlx-c declares. Because
-# we now pull mlx in first, we are also responsible for the option defaults that
-# mlx-c would normally set for us.
+# so that our pinned version wins over the one mlx-c declares. Because we now
+# pull mlx in first, we are also responsible for the option defaults that mlx-c
+# would normally set for us.
 set(MLX_BUILD_TESTS OFF)
 set(MLX_BUILD_EXAMPLES OFF)
 set(MLX_BUILD_BENCHMARKS OFF)
@@ -61,16 +49,15 @@ set(MLX_BUILD_PYTHON_BINDINGS OFF)
 FetchContent_Declare(
   mlx
   GIT_REPOSITORY "https://github.com/ml-explore/mlx.git"
-  GIT_TAG "v0.32.2"
-  PATCH_COMMAND ${mlx_patch})
+  GIT_TAG "v0.32.4")
 FetchContent_MakeAvailable(mlx)
 
 # mlx-c
 set(MLX_C_BUILD_EXAMPLES OFF)
 FetchContent_Declare(
   mlx-c
   GIT_REPOSITORY "https://github.com/ml-explore/mlx-c.git"
-  GIT_TAG "ebc88f10caa1b625e6b581437a8dea6df8a70085")
+  GIT_TAG "v0.8.0")
 FetchContent_MakeAvailable(mlx-c)
 
 # swift-numerics

Package.swift

@@ -21,6 +21,7 @@ let noMetalCmlxExcludes = [
     "mlx/mlx/backend/metal/compiled.cpp",
     "mlx/mlx/backend/metal/conv.cpp",
     "mlx/mlx/backend/metal/copy.cpp",
+    "mlx/mlx/backend/metal/cross_entropy.cpp",
     "mlx/mlx/backend/metal/custom_kernel.cpp",
     "mlx/mlx/backend/metal/device.cpp",
     "mlx/mlx/backend/metal/device_info.cpp",
@@ -29,6 +30,7 @@ let noMetalCmlxExcludes = [
     "mlx/mlx/backend/metal/event.cpp",
     "mlx/mlx/backend/metal/fence.cpp",
     "mlx/mlx/backend/metal/fft.cpp",
+    "mlx/mlx/backend/metal/gated_delta_update.cpp",
     "mlx/mlx/backend/metal/hadamard.cpp",
     "mlx/mlx/backend/metal/indexing.cpp",
     "mlx/mlx/backend/metal/jit_kernels.cpp",
@@ -245,6 +247,7 @@ let cmlx = Target.target(
         .headerSearchPath("json/single_include/nlohmann"),
         .headerSearchPath("fmt/include"),
         .define("MLX_VERSION", to: "\"0.32.0\""),
+        .unsafeFlags(["-Wno-shorten-64-to-32"]),
     ],
     linkerSettings: linkerSettings,
     plugins: cudaBuildPlugins,

Plugins/CudaBuild/plugin.swift

@@ -16,6 +16,7 @@ struct CudaBuild: BuildToolPlugin {
         var cppLanguageStandard: String? = nil
         var verbose: Bool = false
         var codeGeneration: [CodeGeneration] = []
+        var hostCompiler: String? = nil
 
         init() {}
 
@@ -28,9 +29,29 @@ struct CudaBuild: BuildToolPlugin {
             verbose = try c.decodeIfPresent(Bool.self, forKey: .verbose) ?? false
             codeGeneration =
                 try c.decodeIfPresent([CodeGeneration].self, forKey: .codeGeneration) ?? []
+            hostCompiler = try c.decodeIfPresent(String.self, forKey: .hostCompiler)
         }
     }
 
+    /// Host compiler nvcc should use, in priority order: the
+    /// `MLX_CUDA_HOST_COMPILER` environment variable, the `hostCompiler` key in
+    /// CudaBuild.json, then the toolchain's own clang++.
+    ///
+    /// It must still be a clang: nvcc preprocesses with it and SwiftPM compiles
+    /// the generated `.cpp`, so a GCC-preprocessed translation unit will not build.
+    func resolveHostCompiler(context: PluginContext, settings: Settings) -> String {
+        if let env = ProcessInfo.processInfo.environment["MLX_CUDA_HOST_COMPILER"], !env.isEmpty {
+            return env
+        }
+        if let configured = settings.hostCompiler, !configured.isEmpty {
+            return configured
+        }
+        guard let clangUrl = try? context.tool(named: "clang++") else {
+            fatalError("clang++ not found")
+        }
+        return clangUrl.url.path
+    }
+
     func createBuildCommands(context: PluginContext, target: Target) async throws -> [Command] {
 
         print("CUDA Build Plugin")
@@ -40,12 +61,6 @@ struct CudaBuild: BuildToolPlugin {
             return []
         }
 
-        guard let clangUrl = try? context.tool(named: "clang++") else {
-            fatalError("clang++ not found")
-        }
-
-        print("Use clang++ at: \(clangUrl.url.path)")
-
         let sourceDir = target.directoryURL
 
         print("Source directory: \(sourceDir.path)")
@@ -61,6 +76,10 @@ struct CudaBuild: BuildToolPlugin {
             settings = Settings()
         }
 
+        let hostCompiler = resolveHostCompiler(context: context, settings: settings)
+
+        print("Use clang++ at: \(hostCompiler)")
+
         // Scan source directory for .cu files
 
         let sourceDirPath = sourceDir.path.hasSuffix("/") ? sourceDir.path : sourceDir.path + "/"
@@ -132,6 +151,22 @@ struct CudaBuild: BuildToolPlugin {
         }
         let stdArgs = settings.cppLanguageStandard.map { ["--std", $0] } ?? []
 
+        // Include roots that live outside the package: the CCCL copy to compile
+        // device code against, plus anything named in MLX_CUDA_INCLUDE_PATHS
+        // (cudnn-frontend, CUTLASS). Kept in step with Package.swift, which
+        // applies the same roots to the C++ sources.
+        let environment = ProcessInfo.processInfo.environment
+        var externalIncludeRoots: [String] = []
+        if let ccclDir = environment["MLX_CCCL_DIR"] {
+            externalIncludeRoots.append(ccclDir)
+        }
+        let extraRoots = environment["MLX_CUDA_INCLUDE_PATHS"] ?? ""
+        externalIncludeRoots += extraRoots.split(separator: ":").map(String.init)
+        let externalIncludeArgs: [String] =
+            externalIncludeRoots
+            .filter { !$0.isEmpty }
+            .flatMap { ["-I", $0] }
+
         for inputFile in sourceCuFiles + generatedCuFiles {
             let outputCpp = URL(string: inputFile.relativePath, relativeTo: outputDir)!
                 .deletingPathExtension().appendingPathExtension("cpp")
@@ -141,12 +176,13 @@ struct CudaBuild: BuildToolPlugin {
                         "Compiling \(inputFile.lastPathComponent) to \(outputCpp.lastPathComponent)",
                     executable: encuda.url,
                     arguments: compileArguments(
-                        clangpp: clangUrl.url.path,
+                        clangpp: hostCompiler,
                         sourceDir: sourceDir.path,
                         verboseFlag: verboseFlag,
                         stdArgs: stdArgs,
                         incrementalFlag: incrementalFlag,
                         headerSearchPathArgs: headerSearchPathArgs,
+                        externalIncludeArgs: externalIncludeArgs,
                         input: inputFile.path,
                         output: outputCpp.path
                     ),
@@ -171,7 +207,7 @@ struct CudaBuild: BuildToolPlugin {
                 displayName: "Linking CUDA objects",
                 executable: encuda.url,
                 arguments: linkArguments(
-                    clangpp: clangUrl.url.path,
+                    clangpp: hostCompiler,
                     verboseFlag: verboseFlag,
                     stdArgs: stdArgs,
                     incrementalFlag: incrementalFlag,
@@ -195,6 +231,7 @@ struct CudaBuild: BuildToolPlugin {
         stdArgs: [String],
         incrementalFlag: [String],
         headerSearchPathArgs: [String],
+        externalIncludeArgs: [String],
         input: String,
         output: String
     ) -> [String] {
@@ -205,6 +242,7 @@ struct CudaBuild: BuildToolPlugin {
         arguments += ["--clangpp", clangpp]
         arguments += ["-I", sourceDir]
         arguments += headerSearchPathArgs
+        arguments += externalIncludeArgs
         arguments += [input]
         arguments += ["-o", output]
         return arguments

Source/Cmlx/include-framework/Metal.hpp

@@ -2,7 +2,7 @@
 //
 // Metal.hpp.in
 //
-// Autogenerated on October 05, 2026.
+// Autogenerated from commit d5839e5e317d5df0f18951e89f0c9c36c8304a73.
 //
 // Copyright 2020-2024 Apple Inc.
 //

Source/Cmlx/include-framework/mlx-allocator.h

@@ -4,9 +4,20 @@
 #pragma once
 
 #include <cstdlib>
+#include <functional>
+#include <utility>
 
 #include <Cmlx/mlx-api.h>
 
+namespace mlx::core::allocator {
+class Data;
+} // namespace mlx::core::allocator
+
+namespace mlx::core::cu {
+class CommandEncoder;
+allocator::Data malloc_async(size_t size, CommandEncoder& encoder);
+} // namespace mlx::core::cu
+
 namespace mlx::core::allocator {
 
 // Simple wrapper around buffer pointers
@@ -52,12 +63,47 @@ class MLX_API Allocator {
 
 MLX_API Allocator& allocator();
 
-inline Buffer malloc(size_t size) {
-  return allocator().malloc(size);
-}
+using Deleter = std::function<void(Buffer)>;
+
+// The allocator owns data from malloc
+class Data {
+ public:
+  Data(Buffer buffer, Deleter d) : buffer_(buffer), deleter_(std::move(d)) {}
+  Data(const Data& other) = delete;
+  Data& operator=(const Data& other) = delete;
+  Data(Data&& other) noexcept
+      : buffer_(std::exchange(other.buffer_, Buffer(nullptr))),
+        deleter_(std::exchange(other.deleter_, nullptr)),
+        owned_(std::exchange(other.owned_, false)) {}
+  ~Data() {
+    if (owned_) {
+      allocator().free(buffer_);
+    } else if (deleter_) {
+      deleter_(buffer_);
+    }
+  }
+
+  Buffer buffer() const& {
+    return buffer_;
+  }
+  Buffer buffer() const&& = delete;
+
+  bool is_owned() const {
+    return owned_;
+  }
+
+ private:
+  Buffer buffer_;
+  Deleter deleter_;
+  bool owned_{false};
+
+  explicit Data(Buffer buffer) : buffer_(buffer), owned_(true) {}
+  friend Data malloc(size_t size);
+  friend Data cu::malloc_async(size_t size, cu::CommandEncoder& encoder);
+};
 
-inline void free(Buffer buffer) {
-  allocator().free(buffer);
+inline Data malloc(size_t size) {
+  return Data(allocator().malloc(size));
 }
 
 // Make a Buffer from a raw pointer of the given size without a copy.  If a

Source/Cmlx/include-framework/mlx-array.h

@@ -20,7 +20,6 @@ namespace mlx::core {
 // Forward declaration
 class Primitive;
 
-using Deleter = std::function<void(allocator::Buffer)>;
 using ShapeElem = int32_t;
 using Shape = SmallVector<ShapeElem>;
 using Strides = SmallVector<int64_t>;
@@ -70,12 +69,8 @@ class MLX_API array {
       Dtype dtype,
       const std::function<void(void*)>& deleter);
 
-  /* Build an array from a buffer */
-  explicit array(
-      allocator::Buffer data,
-      Shape shape,
-      Dtype dtype,
-      Deleter deleter = allocator::free);
+  /* Build an array from data */
+  explicit array(allocator::Data data, Shape shape, Dtype dtype);
 
   /** Assignment to rvalue does not compile. */
   array& operator=(const array& other) && = delete;
@@ -239,22 +234,7 @@ class MLX_API array {
     return reinterpret_cast<std::uintptr_t>(array_desc_->primitive.get());
   }
 
-  struct Data {
-    allocator::Buffer buffer;
-    Deleter d;
-    Data(allocator::Buffer buffer, Deleter d = allocator::free)
-        : buffer(buffer), d(d) {}
-    // Not copyable
-    Data(const Data& d) = delete;
-    Data& operator=(const Data& d) = delete;
-    Data(Data&& o) : buffer(o.buffer), d(o.d) {
-      o.buffer = allocator::Buffer(nullptr);
-      o.d = [](allocator::Buffer) {};
-    }
-    ~Data() {
-      d(buffer);
-    }
-  };
+  using Data = allocator::Data;
 
   struct Flags {
     // True iff there are no gaps in the underlying data. Each item
@@ -303,7 +283,8 @@ class MLX_API array {
 
   /** True indicates the arrays buffer is safe to reuse */
   bool is_donatable() const {
-    return array_desc_.use_count() == 1 && (array_desc_->data.use_count() == 1);
+    return array_desc_.use_count() == 1 &&
+        (array_desc_->data.use_count() == 1) && array_desc_->data->is_owned();
   }
 
   /** The array's siblings. */
@@ -360,11 +341,8 @@ class MLX_API array {
     return array_desc_->data_size;
   }
 
-  allocator::Buffer& buffer() {
-    return array_desc_->data->buffer;
-  }
-  const allocator::Buffer& buffer() const {
-    return array_desc_->data->buffer;
+  allocator::Buffer buffer() const {
+    return array_desc_->data->buffer();
   }
 
   size_t buffer_size() const {
@@ -448,24 +426,14 @@ class MLX_API array {
   // Check if the array is a tracer array
   bool is_tracer() const;
 
-  void set_data(allocator::Buffer buffer, Deleter d = allocator::free);
+  void set_data(allocator::Data data);
 
   void set_data(
-      allocator::Buffer buffer,
+      allocator::Data data,
       size_t data_size,
       Strides strides,
       Flags flags,
-      Deleter d) {
-    set_data(buffer, data_size, std::move(strides), flags, 0, std::move(d));
-  }
-
-  void set_data(
-      allocator::Buffer buffer,
-      size_t data_size,
-      Strides strides,
-      Flags flags,
-      int64_t offset = 0,
-      Deleter d = allocator::free);
+      int64_t offset = 0);
 
   void copy_shared_buffer(
       const array& other,