Merge remote-tracking branch 'origin/branch-22.06' into feature/skbuild

# Conflicts:
#	build.sh
#	conda/recipes/pylibraft/build.sh

build.sh

@@ -18,8 +18,8 @@ ARGS=$*
 # script, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libraft pyraft pylibraft docs tests bench clean -v -g --install --compile-libs --compile-nn --compile-dist --allgpuarch --nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
-HELP="$0 [<target> ...] [<flag> ...]
+VALIDARGS="clean libraft pyraft pylibraft docs tests bench clean -v -g --install --compile-libs --compile-nn --compile-dist --allgpuarch --no-nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
+HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
    libraft          - build the raft C++ code only. Also builds the C-wrapper library
@@ -41,7 +41,7 @@ HELP="$0 [<target> ...] [<flag> ...]
    --allgpuarch                - build for all supported GPU architectures
    --buildfaiss                - build faiss statically into raft
    --install                   - install cmake targets
-   --nvtx                      - enable nvtx for profiling support
+   --no-nvtx                   - disable nvtx (profiling markers), but allow enabling it in downstream projects
    --show_depr_warn            - show cmake deprecation warnings
    --cmake-args=\\\"<args>\\\" - pass arbitrary list of CMake configuration options (escape all quotes in argument)
    -h                          - print this text
@@ -71,7 +71,7 @@ ENABLE_thrust_DEPENDENCY=ON
 ENABLE_ucx_DEPENDENCY=OFF
 ENABLE_nccl_DEPENDENCY=OFF
 
-NVTX=OFF
+NVTX=ON
 CLEAN=0
 UNINSTALL=0
 DISABLE_DEPRECATION_WARNINGS=ON
@@ -181,8 +181,8 @@ fi
 if hasArg --buildfaiss; then
     BUILD_STATIC_FAISS=ON
 fi
-if hasArg --nvtx; then
-    NVTX=ON
+if hasArg --no-nvtx; then
+    NVTX=OFF
 fi
 if hasArg --show_depr_warn; then
     DISABLE_DEPRECATION_WARNINGS=OFF
@@ -243,7 +243,7 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs || hasArg tests || has
           -DCMAKE_CUDA_ARCHITECTURES=${RAFT_CMAKE_CUDA_ARCHITECTURES} \
           -DRAFT_COMPILE_LIBRARIES=${COMPILE_LIBRARIES} \
           -DRAFT_ENABLE_NN_DEPENDENCIES=${ENABLE_NN_DEPENDENCIES} \
-          -DNVTX=${NVTX} \
+          -DRAFT_NVTX=${NVTX} \
           -DDISABLE_DEPRECATION_WARNINGS=${DISABLE_DEPRECATION_WARNINGS} \
           -DBUILD_TESTS=${BUILD_TESTS} \
           -DBUILD_BENCH=${BUILD_BENCH} \

conda/environments/raft_dev_cuda11.0.yml

@@ -16,6 +16,7 @@ dependencies:
 - rapids-doc-env=22.02.*
 - rmm=22.02.*
 - dask-cuda=22.02.*
+- ucx>=1.12.1
 - ucx-py=0.26.*
 - ucx-proc=*=gpu
 - doxygen>=1.8.20

conda/environments/raft_dev_cuda11.2.yml

@@ -16,6 +16,7 @@ dependencies:
 - rapids-doc-env=22.02.*
 - rmm=22.02.*
 - dask-cuda=22.02.*
+- ucx>=1.12.1
 - ucx-py=0.26.*
 - ucx-proc=*=gpu
 - doxygen>=1.8.20

conda/environments/raft_dev_cuda11.4.yml

@@ -16,6 +16,7 @@ dependencies:
 - rapids-doc-env=22.02.*
 - rmm=22.02.*
 - dask-cuda=22.02.*
+- ucx>=1.12.1
 - ucx-py=0.26.*
 - ucx-proc=*=gpu
 - doxygen>=1.8.20

conda/environments/raft_dev_cuda11.5.yml

@@ -17,6 +17,7 @@ dependencies:
 - rapids-doc-env=22.02.*
 - rmm=22.02.*
 - dask-cuda=22.02.*
+- ucx>=1.12.1
 - ucx-py=0.26.*
 - ucx-proc=*=gpu
 - doxygen>=1.8.20

conda/recipes/libraft_distance/build.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
-./build.sh libraft --install -v --allgpuarch --compile-dist
+./build.sh libraft --install -v --allgpuarch --compile-dist --no-nvtx

conda/recipes/libraft_headers/build.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
-./build.sh libraft --install -v --allgpuarch
+./build.sh libraft --install -v --allgpuarch --no-nvtx

conda/recipes/libraft_nn/build.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
-./build.sh libraft --install -v --allgpuarch --compile-nn
+./build.sh libraft --install -v --allgpuarch --compile-nn --no-nvtx

conda/recipes/pylibraft/build.sh

@@ -2,4 +2,4 @@
 #!/usr/bin/env bash
 
 # This assumes the script is executed from the root of the repo directory
-./build.sh pylibraft --install  --cmake-args=\"-DFIND_RAFT_CPP=ON\"
+./build.sh pylibraft --install --no-nvtx --cmake-args=\"-DFIND_RAFT_CPP=ON\"

conda/recipes/pyraft/build.sh

@@ -3,4 +3,4 @@
 # Copyright (c) 2022, NVIDIA CORPORATION.
 
 # This assumes the script is executed from the root of the repo directory
-./build.sh pyraft --install
+./build.sh pyraft --install --no-nvtx

conda/recipes/pyraft/conda_build_config.yaml

@@ -0,0 +1,2 @@
+ucx_version:
+  - "1.12.1"

conda/recipes/pyraft/meta.yaml

@@ -34,6 +34,7 @@ requirements:
     - cudatoolkit {{ cuda_version }}.*
     - cuda-python >=11.5,<12.0
     - nccl>=2.9.9
+    - ucx {{ ucx_version }}
     - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
   run:
@@ -42,6 +43,7 @@ requirements:
     - libraft-headers {{ version }}
     - nccl>=2.9.9
     - rmm {{ minor_version }}
+    - ucx >={{ ucx_version }}
     - ucx-py {{ ucx_py_version }}
     - ucx-proc=*=gpu
     - dask>=2022.03.0

cpp/CMakeLists.txt

@@ -59,7 +59,7 @@ option(CUDA_STATIC_RUNTIME "Statically link the CUDA runtime" 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(NVTX "Enable nvtx markers" OFF)
+option(RAFT_NVTX "Enable nvtx markers" OFF)
 
 option(RAFT_COMPILE_LIBRARIES "Enable building raft shared library instantiations" ${BUILD_TESTS})
 option(RAFT_COMPILE_NN_LIBRARY "Enable building raft nearest neighbors shared library instantiations" ${RAFT_COMPILE_LIBRARIES})
@@ -91,7 +91,7 @@ message(VERBOSE "RAFT: Disable depreaction warnings " ${DISABLE_DEPRECATION_WARN
 message(VERBOSE "RAFT: Disable OpenMP: ${DISABLE_OPENMP}")
 message(VERBOSE "RAFT: Enable kernel resource usage info: ${CUDA_ENABLE_KERNELINFO}")
 message(VERBOSE "RAFT: Enable lineinfo in nvcc: ${CUDA_ENABLE_LINEINFO}")
-message(VERBOSE "RAFT: Enable nvtx markers: ${NVTX}")
+message(VERBOSE "RAFT: Enable nvtx markers: ${RAFT_NVTX}")
 message(VERBOSE "RAFT: Statically link the CUDA runtime: ${CUDA_STATIC_RUNTIME}")
 
 # Set RMM logging level
@@ -180,7 +180,6 @@ target_include_directories(raft INTERFACE
 # Only CUDA libs, rmm, and mdspan should
 # be used in global target.
 target_link_libraries(raft INTERFACE
-        $<$<BOOL:${NVTX}>:CUDA::nvToolsExt>
         CUDA::cublas
         CUDA::curand
         CUDA::cusolver
@@ -191,7 +190,6 @@ target_link_libraries(raft INTERFACE
         $<$<BOOL:${RAFT_ENABLE_mdspan_DEPENDENCY}>:std::mdspan>
 )
 
-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_DIST_LIBRARY OR RAFT_COMPILE_NN_LIBRARY)
@@ -205,6 +203,28 @@ SECTIONS
 ]=])
 endif()
 
+##############################################################################
+# - NVTX support in raft -----------------------------------------------------
+
+if (RAFT_NVTX)
+  # This enables NVTX within the project with no option to disable it downstream.
+  target_link_libraries(raft INTERFACE CUDA::nvToolsExt)
+  target_compile_definitions(raft INTERFACE NVTX_ENABLED)
+else()
+  # Allow enable NVTX downstream if not set here.
+  # This creates a new option at build/install time, which is set by default to OFF,
+  # but can be enabled in the dependent project.
+  get_property(nvtx_option_help_string CACHE RAFT_NVTX PROPERTY HELPSTRING)
+  string(CONCAT nvtx_export_string
+  "option(RAFT_NVTX \"" ${nvtx_option_help_string} "\" OFF)"
+  [=[
+
+target_link_libraries(raft::raft INTERFACE $<$<BOOL:${RAFT_NVTX}>:CUDA::nvToolsExt>)
+target_compile_definitions(raft::raft INTERFACE $<$<BOOL:${RAFT_NVTX}>:NVTX_ENABLED>)
+
+  ]=])
+endif()
+
 ##############################################################################
 # - raft_distance ------------------------------------------------------------
 add_library(raft_distance INTERFACE)
@@ -342,18 +362,23 @@ include(CPack)
 
 install(TARGETS raft
         DESTINATION ${lib_dir}
+        COMPONENT raft
         EXPORT raft-exports)
 
 install(TARGETS raft_distance
         DESTINATION ${lib_dir}
+        COMPONENT raft_distance
         EXPORT raft-distance-exports)
+
 install(TARGETS raft_nn
         DESTINATION ${lib_dir}
+        COMPONENT raft_nn
         EXPORT raft-nn-exports)
 
 if(TARGET raft_distance_lib)
   install(TARGETS raft_distance_lib
           DESTINATION ${lib_dir}
+          COMPONENT raft_distance
           EXPORT raft-distance-lib-exports)
   install(DIRECTORY include/raft_distance
           DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
@@ -362,18 +387,22 @@ endif()
 if(TARGET raft_nn_lib)
   install(TARGETS raft_nn_lib
           DESTINATION ${lib_dir}
+          COMPONENT raft_nn
           EXPORT raft-nn-lib-exports)
 endif()
 
 
 install(DIRECTORY include/raft
+        COMPONENT raft
         DESTINATION ${CMAKE_INSTALL_INCLUDEDIR})
 
 # Temporary install of raft.hpp while the file is removed
 install(FILES include/raft.hpp
+        COMPONENT raft
         DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/raft)
 
 install(FILES ${CMAKE_CURRENT_BINARY_DIR}/include/raft/version_config.hpp
+        COMPONENT raft
         DESTINATION include/raft)
 
 ##############################################################################
@@ -396,7 +425,8 @@ Imported Targets:
 
 ]=])
 
-set(code_string )
+set(code_string ${nvtx_export_string})
+
 if(RAFT_ENABLE_thrust_DEPENDENCY)
   string(APPEND code_string
   [=[
@@ -449,19 +479,23 @@ raft_export(BUILD raft
 include("${rapids-cmake-dir}/export/write_dependencies.cmake")
 
 set(raft_components distance nn)
+set(raft_install_comp distance nn)
 if(TARGET raft_distance_lib)
   list(APPEND raft_components distance-lib)
+  list(APPEND raft_install_comp distance)
 endif()
 if(TARGET raft_nn_lib)
   list(APPEND raft_components  nn-lib)
+  list(APPEND raft_install_comp nn)
 endif()
 
-foreach(comp IN LISTS raft_components)
+foreach(comp install_comp IN ZIP_LISTS raft_components raft_install_comp)
   install(
           EXPORT raft-${comp}-exports
           FILE raft-${comp}-targets.cmake
           NAMESPACE raft::
           DESTINATION "${lib_dir}/cmake/raft"
+          COMPONENT raft_${install_comp}
   )
   export(
           EXPORT raft-${comp}-exports
@@ -472,7 +506,7 @@ foreach(comp IN LISTS raft_components)
           BUILD raft-${comp}-exports "${PROJECT_BINARY_DIR}/raft-${comp}-dependencies.cmake"
   )
   rapids_export_write_dependencies(
-          INSTALL raft-${comp}-exports "${PROJECT_BINARY_DIR}/rapids-cmake/raft/export/raft-${comp}-dependencies.cmake"
+          INSTALL raft-${comp}-exports "${PROJECT_BINARY_DIR}/rapids-cmake/raft/export/${install_comp}/raft-${comp}-dependencies.cmake"
   )
 
 endforeach()

cpp/bench/common/benchmark.hpp

@@ -18,6 +18,7 @@
 
 #include <memory>
 
+#include <raft/core/handle.hpp>
 #include <raft/cudart_utils.h>
 #include <raft/interruptible.hpp>
 
@@ -102,13 +103,13 @@ struct cuda_event_timer {
 /** Main fixture to be inherited and used by all other c++ benchmarks */
 class fixture {
  private:
-  rmm::cuda_stream stream_owner_{};
   rmm::device_buffer scratch_buf_;
 
  public:
+  raft::handle_t handle;
   rmm::cuda_stream_view stream;
 
-  fixture() : stream{stream_owner_.view()}
+  fixture() : stream{handle.get_stream()}
   {
     int l2_cache_size = 0;
     int device_id     = 0;

cpp/bench/random/permute.cu

@@ -37,7 +37,7 @@ struct permute : public fixture {
       in(p.rows * p.cols, stream)
   {
     raft::random::RngState r(123456ULL);
-    uniform(r, in.data(), p.rows, T(-1.0), T(1.0), stream);
+    uniform(handle, r, in.data(), p.rows, T(-1.0), T(1.0));
   }
 
   void run_benchmark(::benchmark::State& state) override
@@ -50,6 +50,7 @@ struct permute : public fixture {
   }
 
  private:
+  raft::handle_t handle;
   permute_inputs params;
   rmm::device_uvector<T> out, in;
   rmm::device_uvector<int> perms;

cpp/bench/random/rng.cu

@@ -51,23 +51,23 @@ struct rng : public fixture {
     raft::random::RngState r(123456ULL, params.gtype);
     loop_on_state(state, [this, &r]() {
       switch (params.type) {
-        case RNG_Normal: normal(r, ptr.data(), params.len, params.start, params.end, stream); break;
+        case RNG_Normal: normal(handle, r, ptr.data(), params.len, params.start, params.end); break;
         case RNG_LogNormal:
-          lognormal(r, ptr.data(), params.len, params.start, params.end, stream);
+          lognormal(handle, r, ptr.data(), params.len, params.start, params.end);
           break;
         case RNG_Uniform:
-          uniform(r, ptr.data(), params.len, params.start, params.end, stream);
+          uniform(handle, r, ptr.data(), params.len, params.start, params.end);
           break;
-        case RNG_Gumbel: gumbel(r, ptr.data(), params.len, params.start, params.end, stream); break;
+        case RNG_Gumbel: gumbel(handle, r, ptr.data(), params.len, params.start, params.end); break;
         case RNG_Logistic:
-          logistic(r, ptr.data(), params.len, params.start, params.end, stream);
+          logistic(handle, r, ptr.data(), params.len, params.start, params.end);
           break;
-        case RNG_Exp: exponential(r, ptr.data(), params.len, params.start, stream); break;
-        case RNG_Rayleigh: rayleigh(r, ptr.data(), params.len, params.start, stream); break;
+        case RNG_Exp: exponential(handle, r, ptr.data(), params.len, params.start); break;
+        case RNG_Rayleigh: rayleigh(handle, r, ptr.data(), params.len, params.start); break;
         case RNG_Laplace:
-          laplace(r, ptr.data(), params.len, params.start, params.end, stream);
+          laplace(handle, r, ptr.data(), params.len, params.start, params.end);
           break;
-        case RNG_Fill: fill(r, ptr.data(), params.len, params.start, stream); break;
+        case RNG_Fill: fill(handle, r, ptr.data(), params.len, params.start); break;
       };
     });
   }

cpp/bench/spatial/fused_l2_nn.cu

@@ -46,8 +46,8 @@ struct fused_l2_nn : public fixture {
     raft::handle_t handle{stream};
     raft::random::RngState r(123456ULL);
 
-    uniform(r, x.data(), p.m * p.k, T(-1.0), T(1.0), stream);
-    uniform(r, y.data(), p.n * p.k, T(-1.0), T(1.0), stream);
+    uniform(handle, r, x.data(), p.m * p.k, T(-1.0), T(1.0));
+    uniform(handle, r, y.data(), p.n * p.k, T(-1.0), T(1.0));
     raft::linalg::rowNorm(xn.data(), x.data(), p.k, p.m, raft::linalg::L2Norm, true, stream);
     raft::linalg::rowNorm(yn.data(), y.data(), p.k, p.n, raft::linalg::L2Norm, true, stream);
     raft::distance::initialize<T, cub::KeyValuePair<int, T>, int>(

cpp/bench/spatial/selection.cu

@@ -47,7 +47,7 @@ struct selection : public fixture {
   {
     raft::sparse::iota_fill(in_ids_.data(), IdxT(p.n_inputs), IdxT(p.input_len), stream);
     raft::random::RngState state{42};
-    raft::random::uniform(state, in_dists_.data(), in_dists_.size(), KeyT(-1.0), KeyT(1.0), stream);
+    raft::random::uniform(handle, state, in_dists_.data(), in_dists_.size(), KeyT(-1.0), KeyT(1.0));
   }
 
   void run_benchmark(::benchmark::State& state) override