patch for Cuda builds

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

@@ -16,8 +16,15 @@ rm -rf build
 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
-# TODO dkoski -- disabled for now until clear_streams is available
-# ./example1 --device gpu
-# ./tutorial --device gpu
+./example1 --device gpu
+./tutorial --device gpu
 popd

CMakeLists.txt

@@ -14,7 +14,9 @@ set(CMAKE_EXPORT_COMPILE_COMMANDS ON)
 # ----------------------------- Configuration -----------------------------
 # note: mirrors a subset of MLX options exactly (1:1 mapping)
 
-option(MLX_BUILD_EXAMPLES "Build examples for mlx" ON)
+# note: MLX_BUILD_EXAMPLES is owned by mlx itself (see below), so the Swift
+# examples use their own option name to avoid colliding with mlx's C++ examples
+option(MLX_SWIFT_BUILD_EXAMPLES "Build Swift examples for mlx-swift" ON)
 option(MLX_BUILD_METAL "Build metal backend" ON)
 option(MLX_BUILD_CUDA "Build cuda backend" OFF)
 
@@ -26,21 +28,61 @@ if(POLICY CMP0135)
   cmake_policy(SET CMP0135 NEW)
 endif()
 
+# mlx
+
+# note: patches are applied via cmake/apply-patch.cmake so that they are
+# idempotent (the patch step re-runs on every update step) while still failing
+# the build loudly if a patch does not apply -- a silently unpatched dependency
+# is very hard to diagnose (it typically shows up as a runtime failure).
+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.
+set(MLX_BUILD_TESTS OFF)
+set(MLX_BUILD_EXAMPLES OFF)
+set(MLX_BUILD_BENCHMARKS OFF)
+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})
+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 "c74db5307cc8ce122f48d97ef951b30578674e7f")
 FetchContent_MakeAvailable(mlx-c)
 
 # swift-numerics
-set(swift_numerics_patch git apply
-                         ${CMAKE_CURRENT_SOURCE_DIR}/cmake/swift-numerics.patch)
+set(swift_numerics_patch
+    ${apply_patch}
+    -DPATCH=${CMAKE_CURRENT_SOURCE_DIR}/cmake/swift-numerics.patch
+    ${apply_patch_script})
 FetchContent_Declare(
   swift-numerics
   GIT_REPOSITORY "https://github.com/apple/swift-numerics.git"
   GIT_TAG "1.0.2"
-  PATCH_COMMAND ${swift_numerics_patch} || true)
+  PATCH_COMMAND ${swift_numerics_patch})
 FetchContent_MakeAvailable(swift-numerics)
 
 # MLX package
@@ -107,7 +149,7 @@ add_library(MLXLinalg STATIC ${MLXLinalg-src})
 target_link_libraries(MLXLinalg PRIVATE MLX)
 
 # Examples
-if(MLX_BUILD_EXAMPLES)
+if(MLX_SWIFT_BUILD_EXAMPLES)
   add_executable(example1
                  ${CMAKE_CURRENT_LIST_DIR}/Source/Examples/Example1.swift)
   target_link_libraries(example1 PRIVATE MLX)

cmake/apply-patch.cmake

@@ -0,0 +1,69 @@
+# Idempotent `git apply` for FetchContent PATCH_COMMAND.
+#
+# Usage (from a FetchContent_Declare PATCH_COMMAND):
+#
+# PATCH_COMMAND ${CMAKE_COMMAND} -DREPO=<SOURCE_DIR>
+# -DPATCH=${CMAKE_CURRENT_SOURCE_DIR}/cmake/foo.patch -P
+# ${CMAKE_CURRENT_SOURCE_DIR}/cmake/apply-patch.cmake
+#
+# Unlike `git apply ... || true` this does not swallow real failures: an
+# already-applied patch is a no-op (the patch step re-runs on every update), but
+# a patch that fails to apply is a hard error instead of a silently unpatched
+# dependency.
+
+cmake_minimum_required(VERSION 3.16)
+
+if(NOT DEFINED PATCH OR PATCH STREQUAL "")
+  message(FATAL_ERROR "apply-patch.cmake: -DPATCH=<file> is required")
+endif()
+
+if(NOT DEFINED REPO
+   OR REPO STREQUAL ""
+   OR NOT IS_DIRECTORY "${REPO}")
+  # the patch step runs with the dependency source dir as its working directory
+  set(REPO "${CMAKE_CURRENT_BINARY_DIR}")
+endif()
+
+if(NOT EXISTS "${PATCH}")
+  message(FATAL_ERROR "apply-patch.cmake: no such patch file: ${PATCH}")
+endif()
+
+# note: deliberately not `find_package(Git REQUIRED)` -- that pulls in
+# FindPackageHandleStandardArgs in script mode for no benefit here.
+if(NOT DEFINED GIT_EXECUTABLE OR GIT_EXECUTABLE STREQUAL "")
+  find_program(GIT_EXECUTABLE NAMES git git.exe)
+endif()
+
+if(NOT GIT_EXECUTABLE)
+  message(
+    FATAL_ERROR "apply-patch.cmake: git not found (pass -DGIT_EXECUTABLE=)")
+endif()
+
+message(STATUS "apply-patch: repo=${REPO} patch=${PATCH} git=${GIT_EXECUTABLE}")
+
+# already applied? (the patch step re-runs whenever the update step re-runs)
+execute_process(
+  COMMAND ${GIT_EXECUTABLE} apply --reverse --check "${PATCH}"
+  WORKING_DIRECTORY "${REPO}"
+  RESULT_VARIABLE reverse_check
+  OUTPUT_QUIET ERROR_QUIET)
+
+if(reverse_check EQUAL 0)
+  message(STATUS "apply-patch: already applied: ${PATCH}")
+  return()
+endif()
+
+# note: git's stdout/stderr are intentionally *not* captured so that the real
+# reason a patch failed always reaches the build log, even if this script dies
+# before it can format a message of its own.
+execute_process(
+  COMMAND ${GIT_EXECUTABLE} apply --verbose "${PATCH}"
+  WORKING_DIRECTORY "${REPO}"
+  RESULT_VARIABLE apply_result)
+
+if(NOT apply_result EQUAL 0)
+  message(FATAL_ERROR "apply-patch: failed to apply ${PATCH} in ${REPO} "
+                      "(git apply exit ${apply_result}, see output above)")
+endif()
+
+message(STATUS "apply-patch: applied ${PATCH}")

cmake/mlx.patch

@@ -0,0 +1,34 @@
+diff --git a/mlx/backend/cuda/cuda_utils.h b/mlx/backend/cuda/cuda_utils.h
+index f8a234ee6..3b7a79f27 100644
+--- a/mlx/backend/cuda/cuda_utils.h
++++ b/mlx/backend/cuda/cuda_utils.h
+@@ -30,7 +30,11 @@ class CudaHandle {
+     if (cudaPeekAtLastError() != cudaSuccess) {
+       return;
+     }
+-    reset();
++    // Inline reset() without throwing -- handle shutdown cases.
++    if (handle_ != nullptr) {
++      Destroy(handle_);
++      handle_ = nullptr;
++    }
+   }
+
+   CudaHandle(const CudaHandle&) = delete;
+diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp
+index 472b3d99f..e7d8f2620 100644
+--- a/mlx/backend/cuda/device.cpp
++++ b/mlx/backend/cuda/device.cpp
+@@ -617,8 +617,10 @@ std::unordered_map<int, CommandEncoder>& get_command_encoders() {
+ }
+
+ std::unordered_map<int, CommandEncoder>& get_global_command_encoders() {
+-  static std::unordered_map<int, CommandEncoder> encoders;
+-  return encoders;
++  // encoders are leaked intentionally as they would synchronize on process
++  // shutdown
++  static auto* encoders = new std::unordered_map<int, CommandEncoder>();
++  return *encoders;
+ }
+
+ } // namespace mlx::core::cu