diff --git a/mlx/backend/cuda/CMakeLists.txt b/mlx/backend/cuda/CMakeLists.txt index 5fb5096e..2d00fd4c 100644 --- a/mlx/backend/cuda/CMakeLists.txt +++ b/mlx/backend/cuda/CMakeLists.txt @@ -151,16 +151,46 @@ message(STATUS "CUDA architectures: ${MLX_CUDA_ARCHITECTURES}") set_target_properties(mlx PROPERTIES CUDA_ARCHITECTURES "${MLX_CUDA_ARCHITECTURES}") -if(MLX_BUILD_PYTHON_BINDINGS) - set_property( - TARGET mlx - APPEND - PROPERTY INSTALL_RPATH - # The paths here should match the install_requires in setup.py. - "$ORIGIN/../../nvidia/cublas/lib" - "$ORIGIN/../../nvidia/cuda_nvrtc/lib" - "$ORIGIN/../../nvidia/cudnn/lib" - "$ORIGIN/../../nvidia/nccl/lib") +# Search CUDA libs from installed python packages. +if(WIN32) + # Resolve paths of unfound DLL at runtime. + if(BUILD_SHARED_LIBS) + target_link_libraries(mlx PRIVATE "delayimp.lib") + target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/delayload.cpp) + else() + # For static library the delayload must be compiled into final executables. + target_link_libraries(mlx PUBLIC "delayimp.lib") + target_sources( + mlx PUBLIC $) + endif() + # Get all the CUDA DLLs we could link with. + file( + GLOB CUDA_DLL_NAMES + RELATIVE "${CUDAToolkit_BIN_DIR}/x64" + "${CUDAToolkit_BIN_DIR}/x64/*.dll") + # Delay load CUDA and cuDNN libs. + foreach(CUDA_DLL ${CUDA_DLL_NAMES} ${CUDNN_DLL_NAMES}) + target_link_options(mlx PUBLIC "/DELAYLOAD:${CUDA_DLL}") + endforeach() + # Pass the locations where CUDA DLLs are placed. + if(NOT MLX_LOAD_CUDA_LIBS_FROM_PYTHON) + target_compile_definitions( + mlx PUBLIC MLX_CUDA_BIN_DIR="${CUDAToolkit_BIN_DIR}/x64" + MLX_CUDNN_BIN_DIR="${CUDNN_BIN_DIR}") + endif() +else() + # For POSIX we rely on RPATH to search for CUDA libs. + if(MLX_LOAD_CUDA_LIBS_FROM_PYTHON) + set_property( + TARGET mlx + APPEND + PROPERTY INSTALL_RPATH + # The paths here should match the install_requires in setup.py. + "$ORIGIN/../../nvidia/cublas/lib" + "$ORIGIN/../../nvidia/cuda_nvrtc/lib" + "$ORIGIN/../../nvidia/cudnn/lib" + "$ORIGIN/../../nvidia/nccl/lib") + endif() endif() # ------------------------ Dependencies ------------------------ diff --git a/mlx/backend/cuda/delayload.cpp b/mlx/backend/cuda/delayload.cpp new file mode 100644 index 00000000..0ac1cc5e --- /dev/null +++ b/mlx/backend/cuda/delayload.cpp @@ -0,0 +1,80 @@ +// Copyright © 2026 Apple Inc. + +#include "mlx/backend/common/utils.h" + +// clang-format off +#include // must be included first +#include +// clang-format on + +namespace mlx::core { + +namespace fs = std::filesystem; + +inline fs::path relative_to_current_binary(const char* relative) { + return fs::absolute(current_binary_dir() / relative); +} + +inline fs::path cublas_bin_dir() { +#if defined(MLX_CUDA_BIN_DIR) + return MLX_CUDA_BIN_DIR; +#else + return relative_to_current_binary("../nvidia/cublas/bin"); +#endif +} + +fs::path load_nvrtc() { +#if defined(MLX_CUDA_BIN_DIR) + fs::path nvrtc_bin_dir = MLX_CUDA_BIN_DIR; +#else + fs::path nvrtc_bin_dir = + relative_to_current_binary("../nvidia/cuda_nvrtc/bin"); +#endif + // Internally nvrtc loads some libs dynamically, add to search dirs. + ::AddDllDirectory(nvrtc_bin_dir.c_str()); + return nvrtc_bin_dir; +} + +fs::path load_cudnn() { +#if defined(MLX_CUDNN_BIN_DIR) + fs::path cudnn_bin_dir = MLX_CUDNN_BIN_DIR; +#else + fs::path cudnn_bin_dir = relative_to_current_binary("../nvidia/cudnn/bin"); +#endif + // Must load cudnn_graph64_9.dll before locating symbols, otherwise We would + // get errors like "Invalid handle. Cannot load symbol cudnnCreate". + for (const auto& dll : fs::directory_iterator(cudnn_bin_dir)) { + if (dll.path().filename().string().starts_with("cudnn_graph") && + dll.path().extension() == ".dll") { + ::LoadLibraryW(dll.path().c_str()); + break; + } + } + // Internally cuDNN loads some libs dynamically, add to search dirs. + load_nvrtc(); + ::AddDllDirectory(cudnn_bin_dir.c_str()); + ::AddDllDirectory(cublas_bin_dir().c_str()); + return cudnn_bin_dir; +} + +// Called by system when failed to locate a lazy-loaded DLL. +FARPROC WINAPI delayload_helper(unsigned dliNotify, PDelayLoadInfo pdli) { + HMODULE mod = NULL; + if (dliNotify == dliNotePreLoadLibrary) { + std::string dll = pdli->szDll; + if (dll.starts_with("cudnn")) { + static auto cudnn_bin_dir = load_cudnn(); + mod = ::LoadLibraryW((cudnn_bin_dir / dll).c_str()); + } else if (dll.starts_with("cublas")) { + mod = ::LoadLibraryW((cublas_bin_dir() / dll).c_str()); + } else if (dll.starts_with("nvrtc")) { + static auto nvrtc_bin_dir = load_nvrtc(); + mod = ::LoadLibraryW((nvrtc_bin_dir / dll).c_str()); + } + } + return reinterpret_cast(mod); +} + +} // namespace mlx::core + +extern "C" const PfnDliHook __pfnDliNotifyHook2 = mlx::core::delayload_helper; diff --git a/mlx/backend/cuda/jit_module.cpp b/mlx/backend/cuda/jit_module.cpp index 5f9b7362..b0ebb401 100644 --- a/mlx/backend/cuda/jit_module.cpp +++ b/mlx/backend/cuda/jit_module.cpp @@ -31,7 +31,10 @@ const std::vector& include_path_args() { static std::vector cached_args = []() { std::vector args; // Add path to bundled CCCL headers. - auto root_dir = current_binary_dir().parent_path(); + auto root_dir = current_binary_dir(); +#if !defined(_WIN32) + root_dir = root_dir.parent_path(); +#endif auto path = root_dir / "include" / "cccl"; #if defined(MLX_CCCL_DIR) if (!std::filesystem::exists(path)) { diff --git a/mlx/backend/cuda/sort.cu b/mlx/backend/cuda/sort.cu index 2c07be48..c258c013 100644 --- a/mlx/backend/cuda/sort.cu +++ b/mlx/backend/cuda/sort.cu @@ -73,8 +73,8 @@ struct LessThan { __device__ __forceinline__ bool operator()(T a, T b) const { if constexpr (std::is_floating_point_v) { - bool an = std::isnan(a); - bool bn = std::isnan(b); + bool an = cuda::std::isnan(a); + bool bn = cuda::std::isnan(b); if (an | bn) { return (!an) & bn; } diff --git a/setup.py b/setup.py index 6816817d..12505bd1 100644 --- a/setup.py +++ b/setup.py @@ -110,6 +110,8 @@ class CMakeBuild(build_ext): ) ) cmake_args += [f"-DMLX_CUDA_ARCHITECTURES={cuda_archs}"] + # Search CUDA libs from python packages. + cmake_args += ["-DMLX_LOAD_CUDA_LIBS_FROM_PYTHON=ON"] # Some generators require explcitly passing config when building. build_args = ["--config", cfg]