Skip to content

Commit

Permalink
Use CUDA math wheels (#4621)
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:
  - James Lamb (https://github.com/jameslamb)
  - Robert Maynard (https://github.com/robertmaynard)
  - Chuck Hastings (https://github.com/ChuckHastings)
  - Bradley Dice (https://github.com/bdice)

URL: #4621
  • Loading branch information
KyleFromNVIDIA authored Aug 23, 2024
1 parent 97d1641 commit 6b6e8c8
Show file tree
Hide file tree
Showing 7 changed files with 113 additions and 4 deletions.
18 changes: 17 additions & 1 deletion ci/build_wheel.sh
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,23 @@ if [[ ${package_name} == "nx-cugraph" ]] || \
[[ ${package_name} == "cugraph-equivariant" ]]; then
RAPIDS_PY_WHEEL_NAME="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" RAPIDS_PY_WHEEL_PURE="1" rapids-upload-wheels-to-s3 dist
else
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"
)
;;
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="${package_name}_${RAPIDS_PY_CUDA_SUFFIX}" rapids-upload-wheels-to-s3 final_dist
fi
10 changes: 9 additions & 1 deletion ci/build_wheel_cugraph.sh
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,16 @@ export PIP_CONSTRAINT="${PWD}/constraints.txt"

PARALLEL_LEVEL=$(python -c \
"from math import ceil; from multiprocessing import cpu_count; print(ceil(cpu_count()/4))")
case "${RAPIDS_CUDA_VERSION}" in
12.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
;;
11.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF"
;;
esac

export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_CUGRAPH_CPP=OFF;-DCPM_cugraph-ops_SOURCE=${GITHUB_WORKSPACE}/cugraph-ops/"
export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_CUGRAPH_CPP=OFF;-DCPM_cugraph-ops_SOURCE=${GITHUB_WORKSPACE}/cugraph-ops/${EXTRA_CMAKE_ARGS}"
export SKBUILD_BUILD_TOOL_ARGS="-j${PARALLEL_LEVEL};-l${PARALLEL_LEVEL}"

./ci/build_wheel.sh cugraph python/cugraph
11 changes: 10 additions & 1 deletion ci/build_wheel_pylibcugraph.sh
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,16 @@ set -euo pipefail
PARALLEL_LEVEL=$(python -c \
"from math import ceil; from multiprocessing import cpu_count; print(ceil(cpu_count()/4))")

export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_CUGRAPH_CPP=OFF;-DCPM_cugraph-ops_SOURCE=${GITHUB_WORKSPACE}/cugraph-ops/"
case "${RAPIDS_CUDA_VERSION}" in
12.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=ON"
;;
11.*)
EXTRA_CMAKE_ARGS=";-DUSE_CUDA_MATH_WHEELS=OFF"
;;
esac

export SKBUILD_CMAKE_ARGS="-DDETECT_CONDA_ENV=OFF;-DFIND_CUGRAPH_CPP=OFF;-DCPM_cugraph-ops_SOURCE=${GITHUB_WORKSPACE}/cugraph-ops/${EXTRA_CMAKE_ARGS}"
export SKBUILD_BUILD_TOOL_ARGS="-j${PARALLEL_LEVEL};-l${PARALLEL_LEVEL}"

./ci/build_wheel.sh pylibcugraph python/pylibcugraph
31 changes: 31 additions & 0 deletions dependencies.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ files:
extras:
table: project
includes:
- cuda_wheels
- depends_on_rmm
- depends_on_pylibraft
py_test_pylibcugraph:
Expand Down Expand Up @@ -383,6 +384,36 @@ dependencies:
packages:
- cudatoolkit
- cuda-nvtx
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
common_build:
common:
- output_types: [conda, pyproject]
Expand Down
21 changes: 21 additions & 0 deletions python/cugraph/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ option(FIND_CUGRAPH_CPP "Search for existing CUGRAPH C++ installations before de
OFF
)
option(USE_CUGRAPH_OPS "Enable all functions that call cugraph-ops" ON)
option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF)

if(NOT USE_CUGRAPH_OPS)
message(STATUS "Disabling libcugraph functions that reference cugraph-ops")
Expand All @@ -49,6 +50,8 @@ endif()
include(rapids-cython-core)

if(NOT cugraph_FOUND)
find_package(CUDAToolkit REQUIRED)

set(BUILD_TESTS OFF)
set(BUILD_CUGRAPH_MG_TESTS OFF)
set(BUILD_CUGRAPH_OPS_CPP_TESTS OFF)
Expand All @@ -60,8 +63,26 @@ if(NOT cugraph_FOUND)
set(CUGRAPH_EXCLUDE_CUGRAPH_OPS_FROM_ALL ON)
set(ALLOW_CLONE_CUGRAPH_OPS 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 cugraph-cpp EXCLUDE_FROM_ALL)

if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)
set(rpaths
"$ORIGIN/../nvidia/cublas/lib"
"$ORIGIN/../nvidia/cufft/lib"
"$ORIGIN/../nvidia/curand/lib"
"$ORIGIN/../nvidia/cusolver/lib"
"$ORIGIN/../nvidia/cusparse/lib"
"$ORIGIN/../nvidia/nvjitlink/lib"
)
set_property(TARGET cugraph PROPERTY INSTALL_RPATH ${rpaths} APPEND)
endif()

set(cython_lib_dir cugraph)
install(TARGETS cugraph DESTINATION ${cython_lib_dir})
endif()
Expand Down
20 changes: 20 additions & 0 deletions python/pylibcugraph/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@ option(FIND_CUGRAPH_CPP "Search for existing CUGRAPH C++ installations before de
OFF
)
option(USE_CUGRAPH_OPS "Enable all functions that call cugraph-ops" ON)
option(USE_CUDA_MATH_WHEELS "Use the CUDA math wheels instead of the system libraries" OFF)

if(NOT USE_CUGRAPH_OPS)
message(STATUS "Disabling libcugraph functions that reference cugraph-ops")
Expand All @@ -49,6 +50,8 @@ endif()
include(rapids-cython-core)

if (NOT cugraph_FOUND)
find_package(CUDAToolkit REQUIRED)

set(BUILD_TESTS OFF)
set(BUILD_CUGRAPH_MG_TESTS OFF)
set(BUILD_CUGRAPH_OPS_CPP_TESTS OFF)
Expand All @@ -60,8 +63,25 @@ if (NOT cugraph_FOUND)
set(CUGRAPH_EXCLUDE_CUGRAPH_OPS_FROM_ALL ON)
set(ALLOW_CLONE_CUGRAPH_OPS 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 cugraph-cpp EXCLUDE_FROM_ALL)

if(NOT CUDA_STATIC_MATH_LIBRARIES AND USE_CUDA_MATH_WHEELS)
set(rpaths
"$ORIGIN/../nvidia/cublas/lib"
"$ORIGIN/../nvidia/curand/lib"
"$ORIGIN/../nvidia/cusolver/lib"
"$ORIGIN/../nvidia/cusparse/lib"
"$ORIGIN/../nvidia/nvjitlink/lib"
)
set_property(TARGET cugraph PROPERTY INSTALL_RPATH ${rpaths} APPEND)
endif()

set(cython_lib_dir pylibcugraph)
install(TARGETS cugraph DESTINATION ${cython_lib_dir})
install(TARGETS cugraph_c DESTINATION ${cython_lib_dir})
Expand Down
6 changes: 5 additions & 1 deletion python/pylibcugraph/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,10 @@ authors = [
license = { text = "Apache 2.0" }
requires-python = ">=3.9"
dependencies = [
"nvidia-cublas",
"nvidia-curand",
"nvidia-cusolver",
"nvidia-cusparse",
"pylibraft==24.10.*,>=0.0.0a0",
"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`.
Expand Down Expand Up @@ -73,4 +77,4 @@ requires = [
"pylibraft==24.10.*,>=0.0.0a0",
"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`.
matrix-entry = "cuda_suffixed=true"
matrix-entry = "cuda_suffixed=true;use_cuda_wheels=true"

0 comments on commit 6b6e8c8

Please sign in to comment.