Use CUDA math wheels

ci/build_wheel.sh

@@ -25,7 +25,20 @@ SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DDISABLE_DEPRECATION_WARNINGS=ON;-DC
     --no-deps \
     --disable-pip-version-check
 
+case "${RAPIDS_CUDA_VERSION}" in
+  12.*) EXCLUDE_ARGS=(
+    --exclude "libcublas.so.12"
+    --exclude "libcublasLt.so.12"
+    --exclude "libcufft.so.11"
+    --exclude "libcurand.so.10"
+    --exclude "libcusolver.so.11"
+    --exclude "libcusparse.so.12"
+    --exclude "libnvJitLink.so.12"
+  ) ;;
+  11.*) EXCLUDE_ARGS=() ;;
+esac
+
 mkdir -p final_dist
-python -m auditwheel repair -w final_dist dist/*
+python -m auditwheel repair -w final_dist "${EXCLUDE_ARGS[@]}" dist/*
 
 RAPIDS_PY_WHEEL_NAME="cuml_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist

dependencies.yaml

@@ -88,6 +88,7 @@ files:
     extras:
       table: project
     includes:
+      - cuda
       - py_run
   py_test:
     output: pyproject
@@ -406,6 +407,26 @@ dependencies:
               - *libcusolver114
               - *libcusparse_dev114
               - *libcusparse114
+      - output_types: pyproject
+        matrices:
+          - matrix:
+              cuda: "12.*"
+            packages:
+              - nvidia-cublas-cu12
+              - nvidia-cufft-cu12
+              - nvidia-curand-cu12
+              - nvidia-cusparse-cu12
+              - nvidia-cusolver-cu12
+          - matrix:
+              cuda: "11.*"
+            packages:
+          - matrix:
+            packages:
+              - nvidia-cublas
+              - nvidia-cufft
+              - nvidia-curand
+              - nvidia-cusparse
+              - nvidia-cusolver
   docs:
     common:
       - output_types: [conda, requirements]

python/cuml/CMakeLists.txt

@@ -38,6 +38,7 @@ project(
 option(CUML_UNIVERSAL "Build all cuML Python components." ON)
 option(FIND_CUML_CPP "Search for existing CUML C++ installations before defaulting to local files" OFF)
 option(SINGLEGPU "Disable all mnmg components and comms libraries" OFF)
+option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" ON)
 set(CUML_RAFT_CLONE_ON_PIN OFF)
 
 
@@ -72,8 +73,10 @@ include(rapids-cython-core)
 
 set(CUML_PYTHON_TREELITE_TARGET treelite::treelite)
 
-if(NOT ${CUML_CPU})
+if(NOT CUML_CPU)
   if(NOT cuml_FOUND)
+    find_package(CUDAToolkit REQUIRED)
+
     set(BUILD_CUML_TESTS OFF)
     set(BUILD_PRIMS_TESTS OFF)
     set(BUILD_CUML_C_LIBRARY OFF)
@@ -85,11 +88,19 @@ if(NOT ${CUML_CPU})
 
     # Statically link dependencies if building wheels
     set(CUDA_STATIC_RUNTIME ON)
-    set(CUDA_STATIC_MATH_LIBRARIES ON)
     set(CUML_USE_RAFT_STATIC ON)
     set(CUML_USE_FAISS_STATIC ON)
     set(CUML_USE_TREELITE_STATIC ON)
     set(CUML_USE_CUMLPRIMS_MG_STATIC ON)
+    # Link to the CUDA wheels with shared libraries for CUDA 12+
+    if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.0)
+      set(CUDA_STATIC_MATH_LIBRARIES OFF)
+    else()
+      if(USE_CUDA_MATH_WHEELS)
+        message(FATAL_ERROR "Cannot use CUDA math wheels with CUDA < 12.0")
+      endif()
+      set(CUDA_STATIC_MATH_LIBRARIES ON)
+    endif()
     # Don't install the static libs into wheels
     set(CUML_EXCLUDE_RAFT_FROM_ALL ON)
     set(RAFT_EXCLUDE_FAISS_FROM_ALL ON)
@@ -98,6 +109,16 @@ if(NOT ${CUML_CPU})
 
     add_subdirectory(${CUML_CPP_SRC} cuml-cpp EXCLUDE_FROM_ALL)
 
+    if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)
+      set_property(TARGET ${CUML_CPP_TARGET} PROPERTY INSTALL_RPATH
+        "$ORIGIN/../nvidia/cublas/lib"
+        "$ORIGIN/../nvidia/cufft/lib"
+        "$ORIGIN/../nvidia/curand/lib"
+        "$ORIGIN/../nvidia/cusolver/lib"
+        "$ORIGIN/../nvidia/cusparse/lib"
+      )
+    endif()
+
     set(cython_lib_dir cuml)
     install(TARGETS ${CUML_CPP_TARGET} DESTINATION ${cython_lib_dir})
   endif()

python/cuml/pyproject.toml

@@ -86,6 +86,11 @@ dependencies = [
     "dask-cudf==24.10.*,>=0.0.0a0",
     "joblib>=0.11",
     "numba>=0.57",
+    "nvidia-cublas",
+    "nvidia-cufft",
+    "nvidia-curand",
+    "nvidia-cusolver",
+    "nvidia-cusparse",
     "packaging",
     "pylibraft==24.10.*,>=0.0.0a0",
     "raft-dask==24.10.*,>=0.0.0a0",