Adding fatbin to shared libs and fixing conda paths in cpu build

build.sh

@@ -18,7 +18,7 @@ ARGS=$*
 # script, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libraft pyraft docs -v -g --compile-libs --compile-nn --compile-dist --allgpuarch --nvtx --show_depr_warn -h --nogtest --buildfaiss"
+VALIDARGS="clean libraft pyraft docs -v -g --noinstall --compile-libs --compile-nn --compile-dist --allgpuarch --nvtx --show_depr_warn -h --nogtest --buildfaiss"
 HELP="$0 [<target> ...] [<flag> ...]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
@@ -36,6 +36,7 @@ HELP="$0 [<target> ...] [<flag> ...]
    --allgpuarch     - build for all supported GPU architectures
    --buildfaiss     - build faiss statically into raft
    --nogtest        - do not build google tests for libraft
+   --noinstall     - do not install cmake targets
    --nvtx           - Enable nvtx for profiling support
    --show_depr_warn - show cmake deprecation warnings
    -h               - print this text
@@ -55,18 +56,21 @@ BUILD_ALL_GPU_ARCH=0
 BUILD_TESTS=YES
 BUILD_STATIC_FAISS=OFF
 COMPILE_LIBRARIES=${BUILD_TESTS}
+COMPILE_NN_LIBRARY=OFF
+COMPILE_DIST_LIBRARY=OFF
 ENABLE_NN_DEPENDENCIES=${BUILD_TESTS}
 SINGLEGPU=""
 NVTX=OFF
 CLEAN=0
 DISABLE_DEPRECATION_WARNINGS=ON
 CMAKE_TARGET=""
+INSTALL_TARGET="install"
 
 # Set defaults for vars that may not have been defined externally
 #  FIXME: if INSTALL_PREFIX is not set, check PREFIX, then check
 #         CONDA_PREFIX, but there is no fallback from there!
 INSTALL_PREFIX=${INSTALL_PREFIX:=${PREFIX:=${CONDA_PREFIX}}}
-PARALLEL_LEVEL=${PARALLEL_LEVEL:=""}
+PARALLEL_LEVEL=${PARALLEL_LEVEL:=`nproc`}
 BUILD_ABI=${BUILD_ABI:=ON}
 
 # Default to Ninja if generator is not specified
@@ -76,6 +80,10 @@ function hasArg {
     (( ${NUMARGS} != 0 )) && (echo " ${ARGS} " | grep -q " $1 ")
 }
 
+if hasArg --noinstall; then
+    INSTALL_TARGET=""
+fi
+
 if hasArg -h || hasArg --help; then
     echo "${HELP}"
     exit 0
@@ -93,9 +101,8 @@ fi
 
 # Process flags
 if hasArg -v; then
-    VERBOSE_FLAG=-v
-    CMAKE_LOG_LEVEL="--log-level=VERBOSE"
-    set -x
+    VERBOSE_FLAG="-v"
+    CMAKE_LOG_LEVEL="VERBOSE"
 fi
 if hasArg -g; then
     BUILD_TYPE=Debug
@@ -110,13 +117,17 @@ if hasArg --nogtest; then
     ENABLE_NN_DEPENDENCIES=OFF
 fi
 
+if hasArg --compile-libs; then
+  COMPILE_LIBRARIES=ON
+fi
+
 if hasArg --compile-nn || hasArg --compile-libs; then
     ENABLE_NN_DEPENDENCIES=ON
-    COMPILE_LIBRARIES=ON
+    COMPILE_NN_LIBRARY=ON
     CMAKE_TARGET="raft_nn_lib;${CMAKE_TARGET}"
 fi
 if hasArg --compile-dist || hasArg --compile-libs; then
-    COMPILE_LIBRARIES=ON
+    COMPILE_DIST_LIBRARY=ON
     CMAKE_TARGET="raft_distance_lib;${CMAKE_TARGET}"
 fi
 
@@ -167,31 +178,30 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs; then
         echo "Building for *ALL* supported GPU architectures..."
     fi
 
-    cmake -S ${REPODIR}/cpp -B ${CPP_RAFT_BUILD_DIR} ${CMAKE_LOG_LEVEL} \
+    mkdir -p ${CPP_RAFT_BUILD_DIR}
+    cd ${CPP_RAFT_BUILD_DIR}
+    cmake -S ${REPODIR}/cpp -B ${CPP_RAFT_BUILD_DIR} \
           -DCMAKE_INSTALL_PREFIX=${INSTALL_PREFIX} \
           -DCMAKE_CUDA_ARCHITECTURES=${RAFT_CMAKE_CUDA_ARCHITECTURES} \
           -DRAFT_COMPILE_LIBRARIES=${COMPILE_LIBRARIES} \
           -DRAFT_ENABLE_NN_DEPENDENCIES=${ENABLE_NN_DEPENDENCIES} \
           -DNVTX=${NVTX} \
           -DDISABLE_DEPRECATION_WARNINGS=${DISABLE_DEPRECATION_WARNINGS} \
           -DBUILD_TESTS=${BUILD_TESTS} \
-          -DRAFT_USE_FAISS_STATIC=${BUILD_STATIC_FAISS} \
-          ..
+          -DCMAKE_MESSAGE_LOG_LEVEL=${CMAKE_LOG_LEVEL} \
+          -DRAFT_COMPILE_NN_LIBRARY=${COMPILE_NN_LIBRARY} \
+          -DRAFT_COMPILE_DIST_LIBRARY=${COMPILE_DIST_LIBRARY} \
+          -DRAFT_USE_FAISS_STATIC=${BUILD_STATIC_FAISS}
 
   if (( ${NUMARGS} == 0 )) || hasArg libraft; then
       # Run all c++ targets at once
-      if hasArg --compile-nn || hasArg --compile-dist || hasArg --compile-libs; then
-        if ! hasArg --nogtest; then
-          CMAKE_TARGET="test_raft;${CMAKE_TARGET}"
-        fi
-
-        echo "-- Compiling targets: ${CMAKE_TARGET}"
-        cmake --build  ${CPP_RAFT_BUILD_DIR} -j${PARALLEL_LEVEL} ${VERBOSE_FLAG} --target ${CMAKE_TARGET}
-      else
-        cmake --build  ${CPP_RAFT_BUILD_DIR} -j${PARALLEL_LEVEL} ${VERBOSE_FLAG}
+      if ! hasArg --nogtest; then
+        CMAKE_TARGET="${CMAKE_TARGET};test_raft;"
       fi
-  fi
 
+      echo "-- Compiling targets: ${CMAKE_TARGET}, verbose=${VERBOSE_FLAG}"
+      cmake --build  "${CPP_RAFT_BUILD_DIR}" ${VERBOSE_FLAG} -j${PARALLEL_LEVEL} --target ${CMAKE_TARGET} ${INSTALL_TARGET}
+  fi
 fi
 
 # Build and (optionally) install the cuml Python package

ci/cpu/build.sh

@@ -70,14 +70,15 @@ conda config --set ssl_verify False
 # install. This should eliminate a mismatch between different CUDA versions on
 # cpu vs. gpu builds that is problematic with CUDA 11.5 Enhanced Compat.
 if [ "$BUILD_LIBRAFT" == '1' ]; then
-  BUILD_PYRAFT=1
+  BUILD_RAFT=1
   # If we are doing CUDA + Python builds, libraft package is located at ${CONDA_BLD_DIR}
   CONDA_LOCAL_CHANNEL="${CONDA_BLD_DIR}"
 else
   # If we are doing Python builds only, libraft package is placed here by Project Flash
   CONDA_LOCAL_CHANNEL="ci/artifacts/raft/cpu/.conda-bld/"
 fi
 
+gpuci_mamba_retry install -c conda-forge boa
 
 ###############################################################################
 # BUILD - Conda package builds
@@ -86,26 +87,35 @@ fi
 if [ "$BUILD_LIBRAFT" == '1' ]; then
   gpuci_logger "Building conda packages for libraft-nn, libraft-distance, and libraft-headers"
   if [[ -z "$PROJECT_FLASH" || "$PROJECT_FLASH" == "0" ]]; then
-    gpuci_conda_retry build --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft-nn
-    gpuci_conda_retry build --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft-distance
-    gpuci_conda_retry build --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft-headers
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft_headers
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft_nn
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft_distance
   else
-    gpuci_conda_retry build --no-build-id --croot ${CONDA_BLD_DIR} --dirty --no-remove-work-dir conda/recipes/libraft-nn
-    gpuci_conda_retry build --no-build-id --croot ${CONDA_BLD_DIR} --dirty --no-remove-work-dir conda/recipes/libraft-distance
-    gpuci_conda_retry build --no-build-id --croot ${CONDA_BLD_DIR} --dirty --no-remove-work-dir conda/recipes/libraft-headers
-    mkdir -p ${CONDA_BLD_DIR}/libraft
-    mv ${CONDA_BLD_DIR}/work ${CONDA_BLD_DIR}/libraft/work
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} --dirty --no-remove-work-dir conda/recipes/libraft_headers
+    gpuci_logger "`ls ${CONDA_BLD_DIR}/work`"
+    mkdir -p ${CONDA_BLD_DIR}/libraft_headers/work
+    mv ${CONDA_BLD_DIR}/work ${CONDA_BLD_DIR}/libraft_headers/work
+
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} --dirty --no-remove-work-dir conda/recipes/libraft_nn
+    gpuci_logger "`ls ${CONDA_BLD_DIR}/work`"
+    mkdir -p ${CONDA_BLD_DIR}/libraft_nn/work
+    mv ${CONDA_BLD_DIR}/work ${CONDA_BLD_DIR}/libraft_nn/work
+
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} --dirty --no-remove-work-dir conda/recipes/libraft_distance
+    gpuci_logger "`ls ${CONDA_BLD_DIR}/work`"
+    mkdir -p ${CONDA_BLD_DIR}/libraft_distance/work
+    mv ${CONDA_BLD_DIR}/work ${CONDA_BLD_DIR}/libraft_distance/work
   fi
 else
   gpuci_logger "SKIPPING build of conda packages for libraft-nn, libraft-distance and libraft-headers"
 fi
 
-if [ "$BUILD_raft" == "1" ]; then
+if [ "$BUILD_RAFT" == "1" ]; then
   gpuci_logger "Building conda packages for pyraft"
   if [[ -z "$PROJECT_FLASH" || "$PROJECT_FLASH" == "0" ]]; then
-    gpuci_conda_retry build --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/pyraft --python=$PYTHON
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/pyraft --python=$PYTHON
   else
-    gpuci_conda_retry build --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/pyraft -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON
+    gpuci_conda_retry mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/pyraft -c ${CONDA_LOCAL_CHANNEL} --dirty --no-remove-work-dir --python=$PYTHON
     mkdir -p ${CONDA_BLD_DIR}/pyraft
     mv ${CONDA_BLD_DIR}/work ${CONDA_BLD_DIR}/pyraft/work
   fi
@@ -118,4 +128,4 @@ fi
 ################################################################################
 
 gpuci_logger "Upload conda packages"
-source ci/cpu/upload.sh
+source ci/cpu/upload.sh

ci/cpu/prebuild.sh

@@ -14,7 +14,7 @@
 
 if [[ -z "$PROJECT_FLASH" || "$PROJECT_FLASH" == "0" ]]; then
     #If project flash is not activate, always build both
-    export BUILD_PYRAFT=1
+    export BUILD_RAFT=1
     export BUILD_LIBRAFT=1
 fi
 

ci/cpu/upload.sh

@@ -30,10 +30,10 @@ fi
 
 gpuci_logger "Get conda file output locations"
 
-export LIBRAFT_NN_FILE=`conda build --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft_nn --output`
-export LIBRAFT_DISTANCE_FILE=`conda build --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft_distance --output`
-export LIBRAFT_HEADERS_FILE=`conda build --croot ${CONDA_BLD_DIR} conda/recipes/libraft_headers --output`
-export PYRAFT_FILE=`conda build --croot ${CONDA_BLD_DIR} conda/recipes/pyraft --python=$PYTHON --output`
+export LIBRAFT_NN_FILE=`conda mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft_nn --output`
+export LIBRAFT_DISTANCE_FILE=`conda mambabuild --no-build-id --croot ${CONDA_BLD_DIR} conda/recipes/libraft_distance --output`
+export LIBRAFT_HEADERS_FILE=`conda mambabuild --croot ${CONDA_BLD_DIR} conda/recipes/libraft_headers --output`
+export PYRAFT_FILE=`conda mambabuild --croot ${CONDA_BLD_DIR} conda/recipes/pyraft --python=$PYTHON --output`
 
 ################################################################################
 # UPLOAD - Conda packages
@@ -61,7 +61,7 @@ if [[ "$BUILD_LIBRAFT" == "1" && "$UPLOAD_LIBRAFT" == "1" ]]; then
   gpuci_retry anaconda -t ${MY_UPLOAD_KEY} upload -u ${CONDA_USERNAME:-rapidsai} ${LABEL_OPTION} --skip-existing ${LIBRAFT_HEADERS_FILE} --no-progress
 fi
 
-if [[ "$BUILD_PYRAFT" == "1" ]]; then
+if [[ "$BUILD_RAFT" == "1" ]]; then
   test -e ${PYRAFT_FILE}
   echo "Upload pyraft"
   echo ${PYRAFT_FILE}

conda/recipes/libraft_distance/build.sh

@@ -1,3 +1,3 @@
 #!/usr/bin/env bash
 
-./build.sh clean libraft -v --allgpuarch --compile-dist --nogtest
+./build.sh libraft -v --allgpuarch --compile-dist --nogtest

conda/recipes/libraft_distance/meta.yaml

@@ -36,7 +36,7 @@ requirements:
   build:
     - cmake>=3.20.1
   host:
-    - libraft-headers {{
+    - libraft-headers {{ version }}
     - nccl>=2.9.9
     - cudatoolkit {{ cuda_version }}.*
     - ucx-py {{ ucx_py_version }}

conda/recipes/libraft_headers/build.sh

@@ -1,3 +1,3 @@
 #!/usr/bin/env bash
 
-./build.sh clean libraft -v --allgpuarch --nogtest
+./build.sh libraft -v --allgpuarch --nogtest

conda/recipes/libraft_nn/build.sh

@@ -1,3 +1,3 @@
 #!/usr/bin/env bash
 
-./build.sh clean libraft -v --allgpuarch --compile-nn --nogtest
+./build.sh libraft -v --allgpuarch --compile-nn --nogtest

conda/recipes/pyraft/meta.yaml

@@ -29,16 +29,19 @@ requirements:
     - python x.x
     - setuptools
     - cython>=0.29,<0.30
-    - libraft-headers={{ version }}
+    - rmm {{ minor_version }}
+    - libraft-headers {{ version }}
     - cudatoolkit {{ cuda_version }}.*
+    - cuda-python >=11.5,<12.0
     - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
   run:
     - python x.x
     - dask-cuda {{ minor_version }}
-    - libraft-headers={{ version }}
+    - libraft-headers {{ version }}
     - cupy>=7.8.0,<10.0.0a0
     - nccl>=2.9.9
+    - rmm {{ minor_version }}
     - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
     - dask>=2021.11.1,<=2021.11.2

cpp/CMakeLists.txt

@@ -50,6 +50,8 @@ option(DISABLE_OPENMP "Disable OpenMP" OFF)
 option(NVTX "Enable nvtx markers" OFF)
 
 option(RAFT_COMPILE_LIBRARIES "Enable building raft shared library instantiations" ON)
+option(RAFT_COMPILE_NN_LIBRARY "Enable building raft nearest neighbors shared library instantiations" OFF)
+option(RAFT_COMPILE_DIST_LIBRARY "Enable building raft distant shared library instantiations" OFF)
 option(RAFT_ENABLE_NN_DEPENDENCIES "Search for raft::nn dependencies like faiss" ${RAFT_COMPILE_LIBRARIES})
 include(CMakeDependentOption)
 cmake_dependent_option(RAFT_USE_FAISS_STATIC "Build and statically link the FAISS library for nearest neighbors search on GPU" ON RAFT_COMPILE_LIBRARIES OFF)
@@ -142,6 +144,16 @@ target_link_libraries(raft INTERFACE
 target_compile_definitions(raft INTERFACE $<$<BOOL:${NVTX}>:NVTX_ENABLED>)
 target_compile_features(raft INTERFACE cxx_std_17 $<BUILD_INTERFACE:cuda_std_17>)
 
+if(RAFT_COMPILE_LIBRARIES OR RAFT_COMPILE_DIST_LIBRARY OR RAFT_COMPILE_NN_LIBRARY)
+  file(WRITE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld"
+[=[
+SECTIONS
+{
+.nvFatBinSegment : { *(.nvFatBinSegment) }
+.nv_fatbin : { *(.nv_fatbin) }
+}
+]=])
+endif()
 ##############################################################################
 # - raft_distance ------------------------------------------------------------
 add_library(raft_distance INTERFACE)
@@ -152,7 +164,7 @@ endif()
 
 set_target_properties(raft_distance PROPERTIES EXPORT_NAME distance)
 
-if(RAFT_COMPILE_LIBRARIES)
+if(RAFT_COMPILE_LIBRARIES OR RAFT_COMPILE_DIST_LIBRARY)
   add_library(raft_distance_lib SHARED
     src/distance/specializations/detail
     src/distance/specializations/detail/canberra.cu
@@ -161,14 +173,30 @@ if(RAFT_COMPILE_LIBRARIES)
     src/distance/specializations/detail/cosine.cu
     src/distance/specializations/detail/hamming_unexpanded.cu
     src/distance/specializations/detail/hellinger_expanded.cu
-    src/distance/specializations/detail/jensen_shannon.cu
-    src/distance/specializations/detail/kl_divergence.cu
-    src/distance/specializations/detail/l1.cu
-    src/distance/specializations/detail/l2_expanded.cu
-    src/distance/specializations/detail/l2_sqrt_expanded.cu
-    src/distance/specializations/detail/l2_sqrt_unexpanded.cu
-    src/distance/specializations/detail/l2_unexpanded.cu
-    src/distance/specializations/detail/lp_unexpanded.cu
+    src/distance/specializations/detail/jensen_shannon_float_float_float_int.cu
+    src/distance/specializations/detail/jensen_shannon_float_float_float_uint32.cu
+    src/distance/specializations/detail/jensen_shannon_double_double_double_int.cu
+    src/distance/specializations/detail/kl_divergence_float_float_float_int.cu
+    src/distance/specializations/detail/kl_divergence_float_float_float_uint32.cu
+    src/distance/specializations/detail/kl_divergence_double_double_double_int.cu
+    src/distance/specializations/detail/l1_float_float_float_int.cu
+    src/distance/specializations/detail/l1_float_float_float_uint32.cu
+    src/distance/specializations/detail/l1_double_double_double_int.cu
+    src/distance/specializations/detail/l2_expanded_float_float_float_int.cu
+    src/distance/specializations/detail/l2_expanded_float_float_float_uint32.cu
+    src/distance/specializations/detail/l2_expanded_double_double_double_int.cu
+    src/distance/specializations/detail/l2_sqrt_expanded_float_float_float_int.cu
+    src/distance/specializations/detail/l2_sqrt_expanded_float_float_float_uint32.cu
+    src/distance/specializations/detail/l2_sqrt_expanded_double_double_double_int.cu
+    src/distance/specializations/detail/l2_sqrt_unexpanded_float_float_float_int.cu
+    src/distance/specializations/detail/l2_sqrt_unexpanded_float_float_float_uint32.cu
+    src/distance/specializations/detail/l2_sqrt_unexpanded_double_double_double_int.cu
+    src/distance/specializations/detail/l2_unexpanded_double_double_double_int.cu
+    src/distance/specializations/detail/l2_unexpanded_float_float_float_uint32.cu
+    src/distance/specializations/detail/l2_unexpanded_float_float_float_int.cu
+    src/distance/specializations/detail/lp_unexpanded_double_double_double_int.cu
+    src/distance/specializations/detail/lp_unexpanded_float_float_float_uint32.cu
+    src/distance/specializations/detail/lp_unexpanded_float_float_float_int.cu
   )
   set_target_properties(raft_distance_lib PROPERTIES OUTPUT_NAME raft_distance)
 
@@ -180,6 +208,9 @@ if(RAFT_COMPILE_LIBRARIES)
   target_compile_definitions(raft_distance_lib
           INTERFACE "RAFT_DISTANCE_COMPILED")
 
+  # ensure CUDA symbols aren't relocated to the middle of the debug build binaries
+  target_link_options(raft_distance_lib PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld")
+
 endif()
 
 target_link_libraries(raft_distance INTERFACE raft::raft
@@ -197,7 +228,7 @@ endif()
 
 set_target_properties(raft_nn PROPERTIES EXPORT_NAME nn)
 
-if(RAFT_COMPILE_LIBRARIES)
+if(RAFT_COMPILE_LIBRARIES OR RAFT_COMPILE_NN_LIBRARY)
   add_library(raft_nn_lib SHARED
     src/nn/specializations/ball_cover.cu
     src/nn/specializations/detail/ball_cover_lowdim.cu
@@ -214,6 +245,9 @@ if(RAFT_COMPILE_LIBRARIES)
           PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:${RAFT_CXX_FLAGS}>"
           "$<$<COMPILE_LANGUAGE:CUDA>:${RAFT_CUDA_FLAGS}>"
           )
+  # ensure CUDA symbols aren't relocated to the middle of the debug build binaries
+  target_link_options(raft_nn_lib PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld")
+
   target_compile_definitions(raft_nn_lib
           INTERFACE "RAFT_NN_COMPILED")