FIX RAFT C++ static build for wheels

rapidsai · Mar 3, 2024 · 95002e1 · 95002e1
1 parent 940370d
commit 95002e1
Show file tree

Hide file tree

Showing 4 changed files with 42 additions and 10 deletions.
diff --git a/conda/recipes/cuvs/meta.yaml b/conda/recipes/cuvs/meta.yaml
@@ -46,7 +46,7 @@ requirements:
     {% endif %}
     - cuda-version ={{ cuda_version }}
     - cython >=3.0.0
-    - dlpack >= 0.8
+    - dlpack >=0.8
     - pylibraft {{ minor_version }}
     - libcuvs {{ version }}
     - python x.x

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
@@ -62,12 +62,19 @@ option(CUDA_ENABLE_LINEINFO
        "Enable the -lineinfo option for nvcc (useful for cuda-memcheck / profiler)" OFF
 )
 option(CUDA_STATIC_RUNTIME "Statically link the CUDA toolkit runtime and libraries" OFF)
+option(CUVS_USE_RAFT_STATIC "Build and statically link the RAFT libraries" OFF)
 option(CUDA_LOG_COMPILE_TIME "Write a log of compilation times to nvcc_compile_log.csv" OFF)
 option(DETECT_CONDA_ENV "Enable detection of conda environment for dependencies" ON)
 option(DISABLE_DEPRECATION_WARNINGS "Disable deprecaction warnings " ON)
 option(DISABLE_OPENMP "Disable OpenMP" OFF)
 option(CUVS_NVTX "Enable nvtx markers" OFF)
 
+# The options below allow incorporating libcuvs into another build process
+# without installing all its components. This is useful if total file size is
+# at a premium and we do not expect other consumers to use any APIs of the
+# dependency except those that are directly linked to by the dependent library.
+option(CUVS_EXCLUDE_RAFT_FROM_ALL "Exclude RAFT targets from cuVS's 'all' target" OFF)
+
 if((BUILD_TESTS OR BUILD_C_LIBRARY) AND NOT BUILD_CPU_ONLY)
 
 endif()
@@ -102,6 +109,7 @@ message(VERBOSE "cuVS: Enable nvtx markers: ${CUVS_NVTX}")
 message(VERBOSE
         "cuVS: Statically link the CUDA toolkit runtime and libraries: ${CUDA_STATIC_RUNTIME}"
 )
+message(VERBOSE "cuVS: Build and statically link RAFT libraries: ${CUVS_USE_RAFT_STATIC}")
 
 # Set RMM logging level
 set(RMM_LOGGING_LEVEL

diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake
@@ -21,9 +21,24 @@ function(find_and_configure_raft)
     cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
             "${multiValueArgs}" ${ARGN} )
 
+    if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "branch-${CUML_BRANCH_VERSION_raft}")
+        message(STATUS "cuVS: RAFT pinned tag found: ${PKG_PINNED_TAG}. Cloning raft locally.")
+        set(CPM_DOWNLOAD_raft ON)
+    elseif(PKG_USE_RAFT_STATIC AND (NOT CPM_raft_SOURCE))
+        message(STATUS "cuVS: Cloning raft locally to build static libraries.")
+        set(CPM_DOWNLOAD_raft ON)
+    endif()
+
     set(RAFT_COMPONENTS "")
+
     if(PKG_COMPILE_LIBRARY)
+      if(NOT PKG_USE_RAFT_STATIC)
         string(APPEND RAFT_COMPONENTS " compiled")
+        set(RAFT_COMPILED_LIB raft::compiled PARENT_SCOPE)
+      else()
+        string(APPEND RAFT_COMPONENTS " compiled_static")
+        set(RAFT_COMPILED_LIB raft::compiled_static PARENT_SCOPE)
+      endif()
     endif()
 
     if(PKG_ENABLE_MNMG_DEPENDENCIES)
@@ -39,15 +54,16 @@ function(find_and_configure_raft)
             INSTALL_EXPORT_SET  cuvs-exports
             COMPONENTS          ${RAFT_COMPONENTS}
             CPM_ARGS
-            GIT_REPOSITORY https://github.com/${PKG_FORK}/raft.git
-            GIT_TAG        ${PKG_PINNED_TAG}
-            SOURCE_SUBDIR  cpp
-            OPTIONS
-            "BUILD_TESTS OFF"
-            "BUILD_PRIMS_BENCH OFF"
-            "BUILD_ANN_BENCH OFF"
-            "RAFT_NVTX ${PKG_ENABLE_NVTX}"
-            "RAFT_COMPILE_LIBRARY ${PKG_COMPILE_LIBRARY}"
+              GIT_REPOSITORY        https://github.com/${PKG_FORK}/raft.git
+              GIT_TAG               ${PKG_PINNED_TAG}
+              SOURCE_SUBDIR         cpp
+              EXCLUDE_FROM_ALL      ${PKG_EXCLUDE_FROM_ALL}
+              OPTIONS
+              "BUILD_TESTS OFF"
+              "BUILD_PRIMS_BENCH OFF"
+              "BUILD_ANN_BENCH OFF"
+              "RAFT_NVTX ${PKG_ENABLE_NVTX}"
+              "RAFT_COMPILE_LIBRARY ${PKG_COMPILE_LIBRARY}"
             )
 endfunction()
 
@@ -58,6 +74,8 @@ find_and_configure_raft(VERSION  ${RAFT_VERSION}.00
         FORK                     ${RAFT_FORK}
         PINNED_TAG               ${RAFT_PINNED_TAG}
         COMPILE_LIBRARY          ON
+        USE_RAFT_STATIC          ${CUVS_USE_RAFT_STATIC}
+        EXCLUDE_FROM_ALL         ${CUVS_EXCLUDE_RAFT_FROM_ALL}
         ENABLE_MNMG_DEPENDENCIES OFF
         ENABLE_NVTX              OFF
 )
diff --git a/python/cuvs/CMakeLists.txt b/python/cuvs/CMakeLists.txt
@@ -58,7 +58,13 @@ endif()
 if(NOT cuvs_FOUND)
   set(BUILD_TESTS OFF)
   set(BUILD_C_LIBRARY ON)
+
+  # Statically link dependencies if building wheels
   set(CUDA_STATIC_RUNTIME ON)
+  set(CUVS_USE_RAFT_STATIC ON)
+
+  # Don't install the static libs into wheels
+  set(CUML_EXCLUDE_RAFT_FROM_ALL ON)
 
   add_subdirectory(../../cpp cuvs-cpp EXCLUDE_FROM_ALL)