diff --git a/.github/scripts/build-linux-cuda-cmake.sh b/.github/scripts/build-linux-cuda-cmake.sh index 2a571ea6d..2fc81d589 100755 --- a/.github/scripts/build-linux-cuda-cmake.sh +++ b/.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 diff --git a/CMakeLists.txt b/CMakeLists.txt index 41404a01d..22f4b16ff 100644 --- a/CMakeLists.txt +++ b/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,7 +28,45 @@ 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=) +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" @@ -34,13 +74,15 @@ FetchContent_Declare( 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) diff --git a/cmake/apply-patch.cmake b/cmake/apply-patch.cmake new file mode 100644 index 000000000..63e04f83f --- /dev/null +++ b/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= +# -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= 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}") diff --git a/cmake/mlx.patch b/cmake/mlx.patch new file mode 100644 index 000000000..de99b7ebc --- /dev/null +++ b/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& get_command_encoders() { + } + + std::unordered_map& get_global_command_encoders() { +- static std::unordered_map encoders; +- return encoders; ++ // encoders are leaked intentionally as they would synchronize on process ++ // shutdown ++ static auto* encoders = new std::unordered_map(); ++ return *encoders; + } + + } // namespace mlx::core::cu