Use CUDA math wheels

ci/build_wheel.sh

@@ -20,10 +20,29 @@ rapids-generate-version > VERSION
 
 cd "${package_dir}"
 
+case "${RAPIDS_CUDA_VERSION}" in
+ 12.*)
+ EXCLUDE_ARGS=(
+ --exclude "libcublas.so.12"
+ --exclude "libcublasLt.so.12"
+ --exclude "libcurand.so.10"
+ --exclude "libcusolver.so.11"
+ --exclude "libcusparse.so.12"
+ --exclude "libnvJitLink.so.12"
+ --exclude "libucp.so.0"
+ )
+ ;;
+ 11.*)
+ EXCLUDE_ARGS=(
+ --exclude "libucp.so.0"
+ )
+ ;;
+esac
+
 # Hardcode the output dir
 python -m pip wheel . -w dist -vvv --no-deps --disable-pip-version-check
 
 mkdir -p final_dist
-python -m auditwheel repair -w final_dist --exclude "libucp.so.0" dist/*
+python -m auditwheel repair -w final_dist "${EXCLUDE_ARGS[@]}" dist/*
 
 RAPIDS_PY_WHEEL_NAME="${underscore_package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist

ci/build_wheel_pylibraft.sh

@@ -3,7 +3,16 @@
 
 set -euo pipefail
 
+case "${RAPIDS_CUDA_VERSION}" in
+ 12.*)
+ EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
+ ;;
+ 11.*)
+ EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF"
+ ;;
+esac
+
 # Set up skbuild options. Enable sccache in skbuild config options
-export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF"
+export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_RAFT_CPP=OFF${EXTRA_CMAKE_ARGS}"
 
 ci/build_wheel.sh pylibraft python/pylibraft

dependencies.yaml

@@ -83,6 +83,7 @@ files:
  extras:
  table: project
  includes:
+ - cuda
  - run_pylibraft
  py_test_pylibraft:
  output: pyproject
@@ -377,6 +378,24 @@ dependencies:
  - *libcusolver114
  - *libcusparse_dev114
  - *libcusparse114
+ - output_types: pyproject
+ matrices:
+ - matrix:
+ cuda: "12.*"
+ packages:
+ - nvidia-cublas-cu12
+ - nvidia-curand-cu12
+ - nvidia-cusolver-cu12
+ - nvidia-cusparse-cu12
+ - matrix:
+ cuda: "11.*"
+ packages:
+ - matrix:
+ packages:
+ - nvidia-cublas
+ - nvidia-curand
+ - nvidia-cusolver
+ - nvidia-cusparse
 
  depends_on_cupy:
  common:

python/pylibraft/CMakeLists.txt

@@ -30,6 +30,7 @@ project(
 option(FIND_RAFT_CPP "Search for existing RAFT C++ installations before defaulting to local files"
  ON
 )
+option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF)
 
 # If the user requested it we attempt to find RAFT.
 if(FIND_RAFT_CPP)
@@ -48,15 +49,33 @@ endif()
 include(rapids-cython-core)
 
 if(NOT raft_FOUND)
+ find_package(CUDAToolkit REQUIRED)
+
  set(BUILD_TESTS OFF)
  set(BUILD_PRIMS_BENCH OFF)
  set(BUILD_ANN_BENCH OFF)
  set(RAFT_COMPILE_LIBRARY ON)
  set(CUDA_STATIC_RUNTIME ON)
  set(CUDA_STATIC_MATH_LIBRARIES ON)
+ if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 12.0)
+ set(CUDA_STATIC_MATH_LIBRARIES OFF)
+ elseif(USE_CUDA_MATH_WHEELS)
+ message(FATAL_ERROR "Cannot use CUDA math wheels with CUDA < 12.0")
+ endif()
 
  add_subdirectory(../../cpp raft-cpp EXCLUDE_FROM_ALL)
 
+ if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)
+ set_property(TARGET raft_lib PROPERTY INSTALL_RPATH
+ "$ORIGIN/../nvidia/cublas/lib"
+ "$ORIGIN/../nvidia/cufft/lib"
+ "$ORIGIN/../nvidia/curand/lib"
+ "$ORIGIN/../nvidia/cusolver/lib"
+ "$ORIGIN/../nvidia/cusparse/lib"
+ "$ORIGIN/../nvidia/nvjitlink/lib"
+ )
+ endif()
+
  # When building the C++ libraries from source we must copy libraft.so alongside the
  # pairwise_distance and random Cython libraries TODO: when we have a single 'compiled' raft
  # library, we shouldn't need this

python/pylibraft/pyproject.toml

@@ -33,6 +33,10 @@ requires-python = ">=3.9"
 dependencies = [
  "cuda-python",
  "numpy>=1.23,<2.0a0",
+ "nvidia-cublas",
+ "nvidia-curand",
+ "nvidia-cusolver",
+ "nvidia-cusparse",
  "rmm==24.10.*,>=0.0.0a0",
 ] # This list was generated by `rapids-dependency-file-generator`. To make changes, edit ../../dependencies.yaml and run `rapids-dependency-file-generator`.
 classifiers = [