Skip to content

Commit

Permalink
Use CUDA math wheels (#2415)
Browse files Browse the repository at this point in the history
Use CUDA math wheels to reduce wheel size by not statically linking CUDA math libraries.

Contributes to rapidsai/build-planning#35

Authors:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)

Approvers:
  - Robert Maynard (https://github.com/robertmaynard)
  - Bradley Dice (https://github.com/bdice)
  - James Lamb (https://github.com/jameslamb)

URL: #2415
  • Loading branch information
KyleFromNVIDIA authored Aug 22, 2024
1 parent db07998 commit f1cc0fb
Show file tree
Hide file tree
Showing 5 changed files with 84 additions and 3 deletions.
21 changes: 20 additions & 1 deletion ci/build_wheel.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
11 changes: 10 additions & 1 deletion ci/build_wheel_pylibraft.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
31 changes: 31 additions & 0 deletions dependencies.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -83,6 +83,7 @@ files:
extras:
table: project
includes:
- cuda_wheels
- run_pylibraft
py_test_pylibraft:
output: pyproject
Expand Down Expand Up @@ -377,6 +378,36 @@ dependencies:
- *libcusolver114
- *libcusparse_dev114
- *libcusparse114
cuda_wheels:
specific:
- output_types: pyproject
matrices:
- matrix:
cuda: "12.*"
use_cuda_wheels: "true"
packages:
- nvidia-cublas-cu12
- nvidia-curand-cu12
- nvidia-cusolver-cu12
- nvidia-cusparse-cu12
# CUDA 11 does not provide wheels, so use the system libraries instead
- matrix:
cuda: "11.*"
use_cuda_wheels: "true"
packages:
# if use_cuda_wheels=false is provided, do not add dependencies on any CUDA wheels
# (e.g. for DLFW and pip devcontainers)
- matrix:
use_cuda_wheels: "false"
packages:
# if no matching matrix selectors passed, list the unsuffixed packages
# (just as a source of documentation, as this populates pyproject.toml in source control)
- matrix:
packages:
- nvidia-cublas
- nvidia-curand
- nvidia-cusolver
- nvidia-cusparse

depends_on_cupy:
common:
Expand Down
18 changes: 18 additions & 0 deletions python/pylibraft/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -48,15 +49,32 @@ 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/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
Expand Down
6 changes: 5 additions & 1 deletion python/pylibraft/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down Expand Up @@ -124,7 +128,7 @@ requires = [
"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`.
dependencies-file = "../../dependencies.yaml"
matrix-entry = "cuda_suffixed=true"
matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true"

[tool.pytest.ini_options]
filterwarnings = [
Expand Down

0 comments on commit f1cc0fb

Please sign in to comment.