From 6b6e8c8e0bce8db894ac9dbf6a5e425a6775982b Mon Sep 17 00:00:00 2001 From: Kyle Edwards Date: Fri, 23 Aug 2024 12:14:01 -0400 Subject: [PATCH] Use CUDA math wheels (#4621) Use CUDA math wheels to reduce wheel size by not statically linking CUDA math libraries. Contributes to https://github.com/rapidsai/build-planning/issues/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: https://github.com/rapidsai/cugraph/pull/4621 --- ci/build_wheel.sh | 18 ++++++++++++++++- ci/build_wheel_cugraph.sh | 10 +++++++++- ci/build_wheel_pylibcugraph.sh | 11 ++++++++++- dependencies.yaml | 31 ++++++++++++++++++++++++++++++ python/cugraph/CMakeLists.txt | 21 ++++++++++++++++++++ python/pylibcugraph/CMakeLists.txt | 20 +++++++++++++++++++ python/pylibcugraph/pyproject.toml | 6 +++++- 7 files changed, 113 insertions(+), 4 deletions(-) diff --git a/ci/build_wheel.sh b/ci/build_wheel.sh index 707c61e4d3e..1976d8ff46f 100755 --- a/ci/build_wheel.sh +++ b/ci/build_wheel.sh @@ -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 diff --git a/ci/build_wheel_cugraph.sh b/ci/build_wheel_cugraph.sh index 6f1b23923ff..20d9bf47e3e 100755 --- a/ci/build_wheel_cugraph.sh +++ b/ci/build_wheel_cugraph.sh @@ -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 diff --git a/ci/build_wheel_pylibcugraph.sh b/ci/build_wheel_pylibcugraph.sh index ee33ab4a82d..fa967b0be29 100755 --- a/ci/build_wheel_pylibcugraph.sh +++ b/ci/build_wheel_pylibcugraph.sh @@ -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 diff --git a/dependencies.yaml b/dependencies.yaml index fba62f521e6..42f2fcec63d 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -137,6 +137,7 @@ files: extras: table: project includes: + - cuda_wheels - depends_on_rmm - depends_on_pylibraft py_test_pylibcugraph: @@ -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] diff --git a/python/cugraph/CMakeLists.txt b/python/cugraph/CMakeLists.txt index d96d28b6b59..ca38b5551c9 100644 --- a/python/cugraph/CMakeLists.txt +++ b/python/cugraph/CMakeLists.txt @@ -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") @@ -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) @@ -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() diff --git a/python/pylibcugraph/CMakeLists.txt b/python/pylibcugraph/CMakeLists.txt index d30b03cb68d..045628e9c0d 100644 --- a/python/pylibcugraph/CMakeLists.txt +++ b/python/pylibcugraph/CMakeLists.txt @@ -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") @@ -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) @@ -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}) diff --git a/python/pylibcugraph/pyproject.toml b/python/pylibcugraph/pyproject.toml index 73699b6e445..727d05e8e52 100644 --- a/python/pylibcugraph/pyproject.toml +++ b/python/pylibcugraph/pyproject.toml @@ -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`. @@ -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"