Migrate to use cuVS for vector search

ci/release/update-version.sh

@@ -47,6 +47,7 @@ DEPENDENCIES=(
   libcumlprims
   libraft-headers
   libraft
+  libcuvs
   librmm
   pylibraft
   raft-dask

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -39,6 +39,7 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
+- libcuvs==24.10.*,>=0.0.0a0
 - libraft-headers==24.10.*,>=0.0.0a0
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0

conda/environments/all_cuda-125_arch-x86_64.yaml

@@ -36,6 +36,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libcuvs==24.10.*,>=0.0.0a0
 - libraft-headers==24.10.*,>=0.0.0a0
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0

conda/environments/clang_tidy_cuda-118_arch-x86_64.yaml

@@ -27,6 +27,7 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
+- libcuvs==24.10.*,>=0.0.0a0
 - libraft-headers==24.10.*,>=0.0.0a0
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0

conda/environments/cpp_all_cuda-118_arch-x86_64.yaml

@@ -25,6 +25,7 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
+- libcuvs==24.10.*,>=0.0.0a0
 - libraft-headers==24.10.*,>=0.0.0a0
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0

conda/environments/cpp_all_cuda-125_arch-x86_64.yaml

@@ -22,6 +22,7 @@ dependencies:
 - libcurand-dev
 - libcusolver-dev
 - libcusparse-dev
+- libcuvs==24.10.*,>=0.0.0a0
 - libraft-headers==24.10.*,>=0.0.0a0
 - libraft==24.10.*,>=0.0.0a0
 - librmm==24.10.*,>=0.0.0a0

cpp/CMakeLists.txt

@@ -228,6 +228,7 @@ endif()
 include(cmake/thirdparty/get_cccl.cmake)
 include(cmake/thirdparty/get_rmm.cmake)
 include(cmake/thirdparty/get_raft.cmake)
+include(cmake/thirdparty/get_cuvs.cmake)
 
 if(LINK_TREELITE)
   include(cmake/thirdparty/get_treelite.cmake)
@@ -442,18 +443,6 @@ if(BUILD_CUML_CPP_LIBRARY)
         src/metrics/kl_divergence.cu
         src/metrics/mutual_info_score.cu
         src/metrics/pairwise_distance.cu
-        src/metrics/pairwise_distance_canberra.cu
-        src/metrics/pairwise_distance_chebyshev.cu
-        src/metrics/pairwise_distance_correlation.cu
-        src/metrics/pairwise_distance_cosine.cu
-        src/metrics/pairwise_distance_euclidean.cu
-        src/metrics/pairwise_distance_hamming.cu
-        src/metrics/pairwise_distance_hellinger.cu
-        src/metrics/pairwise_distance_jensen_shannon.cu
-        src/metrics/pairwise_distance_kl_divergence.cu
-        src/metrics/pairwise_distance_l1.cu
-        src/metrics/pairwise_distance_minkowski.cu
-        src/metrics/pairwise_distance_russell_rao.cu
         src/metrics/r2_score.cu
         src/metrics/rand_index.cu
         src/metrics/silhouette_score.cu
@@ -611,6 +600,7 @@ if(BUILD_CUML_CPP_LIBRARY)
   # These are always private:
   list(APPEND _cuml_cpp_private_libs
     raft::raft
+    cuvs::cuvs
     $<TARGET_NAME_IF_EXISTS:GPUTreeShap::GPUTreeShap>
     $<$<BOOL:${LINK_CUFFT}>:CUDA::cufft${_ctk_fft_static_suffix}>
     ${TREELITE_LIBS}

cpp/bench/CMakeLists.txt

@@ -49,6 +49,7 @@ if(BUILD_CUML_BENCH)
       cuml::${CUML_CPP_TARGET}
       benchmark::benchmark
       ${TREELITE_LIBS}
+      cuvs::cuvs
       raft::raft
       raft::compiled
   )

cpp/cmake/thirdparty/get_cuvs.cmake

@@ -0,0 +1,66 @@
+#=============================================================================
+# Copyright (c) 2024, NVIDIA CORPORATION.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#=============================================================================
+
+set(CUML_MIN_VERSION_cuvs "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}.00")
+set(CUML_BRANCH_VERSION_cuvs "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}")
+
+function(find_and_configure_cuvs)
+    set(oneValueArgs VERSION FORK PINNED_TAG EXCLUDE_FROM_ALL COMPILE_LIBRARY CLONE_ON_PIN NVTX)
+    cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
+            "${multiValueArgs}" ${ARGN} )
+
+    if(PKG_CLONE_ON_PIN AND NOT PKG_PINNED_TAG STREQUAL "branch-${CUML_BRANCH_VERSION_cuvs}")
+        message(STATUS "CUML: CUVS pinned tag found: ${PKG_PINNED_TAG}. Cloning cuvs locally.")
+        set(CPM_DOWNLOAD_cuvs ON)
+    endif()
+
+    rapids_cpm_find(cuvs ${PKG_VERSION}
+      GLOBAL_TARGETS      cuvs::cuvs
+      BUILD_EXPORT_SET    cuml-exports
+      INSTALL_EXPORT_SET  cuml-exports
+      CPM_ARGS
+        GIT_REPOSITORY         https://github.com/${PKG_FORK}/cuvs.git
+        GIT_TAG                ${PKG_PINNED_TAG}
+        SOURCE_SUBDIR          cpp
+        EXCLUDE_FROM_ALL       ${PKG_EXCLUDE_FROM_ALL}
+        OPTIONS
+          "BUILD_TESTS OFF"
+          "BUILD_BENCH OFF"
+    )
+
+    if(cuvs_ADDED)
+        message(VERBOSE "CUML: Using CUVS located in ${cuvs_SOURCE_DIR}")
+    else()
+        message(VERBOSE "CUML: Using CUVS located in ${cuvs_DIR}")
+    endif()
+
+
+endfunction()
+
+# Change pinned tag here to test a commit in CI
+# To use a different CUVS locally, set the CMake variable
+# CPM_cuvs_SOURCE=/path/to/local/cuvs
+find_and_configure_cuvs(VERSION          ${CUML_MIN_VERSION_cuvs}
+      FORK             rapidsai
+      PINNED_TAG       branch-${CUML_BRANCH_VERSION_cuvs}
+      EXCLUDE_FROM_ALL ${CUML_EXCLUDE_CUVS_FROM_ALL}
+      # When PINNED_TAG above doesn't match cuml,
+      # force local cuvs clone in build directory
+      # even if it's already installed.
+      CLONE_ON_PIN     ${CUML_CUVS_CLONE_ON_PIN}
+      COMPILE_LIBRARY  ${CUML_CUVS_COMPILED}
+      NVTX             ${NVTX}
+      )

cpp/include/cuml/neighbors/knn.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2019-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2019-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,9 +17,11 @@
 #pragma once
 
 #include <raft/distance/distance_types.hpp>
-#include <raft/spatial/knn/ann_common.h>
 #include <raft/spatial/knn/ball_cover_types.hpp>
 
+#include <cuvs/neighbors/ivf_flat.hpp>
+#include <cuvs/neighbors/ivf_pq.hpp>
+
 namespace raft {
 class handle_t;
 }
@@ -46,6 +48,8 @@ namespace ML {
  *            default
  * @param[in] metric_arg the value of `p` for Minkowski (l-p) distances. This
  *            is ignored if the metric_type is not Minkowski.
+ * @param[in] translations translation ids for indices when index rows represent
+ *        non-contiguous partitions
  */
 void brute_force_knn(const raft::handle_t& handle,
                      std::vector<float*>& input,
@@ -59,7 +63,8 @@ void brute_force_knn(const raft::handle_t& handle,
                      bool rowMajorIndex                  = false,
                      bool rowMajorQuery                  = false,
                      raft::distance::DistanceType metric = raft::distance::DistanceType::L2Expanded,
-                     float metric_arg                    = 2.0f);
+                     float metric_arg                    = 2.0f,
+                     std::vector<int64_t>* translations  = nullptr);
 
 void rbc_build_index(const raft::handle_t& handle,
                      raft::spatial::knn::BallCoverIndex<int64_t, float, uint32_t>& index);
@@ -71,6 +76,35 @@ void rbc_knn_query(const raft::handle_t& handle,
                    uint32_t n_search_items,
                    int64_t* out_inds,
                    float* out_dists);
+
+struct knnIndex {
+  raft::distance::DistanceType metric;
+  float metricArg;
+  int nprobe;
+
+  std::unique_ptr<cuvs::neighbors::ivf_flat::index<float, int64_t>> ivf_flat;
+  std::unique_ptr<cuvs::neighbors::ivf_pq::index<int64_t>> ivf_pq;
+
+  int device;
+};
+
+struct knnIndexParam {
+  virtual ~knnIndexParam() {}
+};
+
+struct IVFParam : knnIndexParam {
+  int nlist;
+  int nprobe;
+};
+
+struct IVFFlatParam : IVFParam {};
+
+struct IVFPQParam : IVFParam {
+  int M;
+  int n_bits;
+  bool usePrecomputedTables;
+};
+
 /**
  * @brief Flat C++ API function to build an approximate nearest neighbors index
  * from an index array and a set of parameters.
@@ -85,8 +119,8 @@ void rbc_knn_query(const raft::handle_t& handle,
  * @param[in] D the dimensionality of the index array
  */
 void approx_knn_build_index(raft::handle_t& handle,
-                            raft::spatial::knn::knnIndex* index,
-                            raft::spatial::knn::knnIndexParam* params,
+                            knnIndex* index,
+                            knnIndexParam* params,
                             raft::distance::DistanceType metric,
                             float metricArg,
                             float* index_array,
@@ -109,7 +143,7 @@ void approx_knn_build_index(raft::handle_t& handle,
 void approx_knn_search(raft::handle_t& handle,
                        float* distances,
                        int64_t* indices,
-                       raft::spatial::knn::knnIndex* index,
+                       knnIndex* index,
                        int k,
                        float* query_array,
                        int n);

cpp/src/hdbscan/detail/soft_clustering.cuh

@@ -24,7 +24,6 @@
 #include <cuml/common/logger.hpp>
 
 #include <raft/core/device_mdspan.hpp>
-#include <raft/distance/distance.cuh>
 #include <raft/distance/distance_types.hpp>
 #include <raft/label/classlabels.cuh>
 #include <raft/linalg/matrix_vector_op.cuh>
@@ -43,6 +42,8 @@
 #include <thrust/execution_policy.h>
 #include <thrust/transform.h>
 
+#include <cuvs/distance/distance.hpp>
+
 #include <algorithm>
 #include <cmath>
 #include <limits>
@@ -88,45 +89,14 @@ void dist_membership_vector(const raft::handle_t& handle,
     value_idx samples_per_batch = min((value_idx)batch_size, (value_idx)n_queries - batch_offset);
     rmm::device_uvector<value_t> dist(samples_per_batch * n_exemplars, stream);
 
-    // compute the distances using raft API
-    switch (metric) {
-      case raft::distance::DistanceType::L2SqrtExpanded:
-        raft::distance::
-          distance<raft::distance::DistanceType::L2SqrtExpanded, value_t, value_t, value_t, int>(
-            handle,
-            query + batch_offset * n,
-            exemplars_dense.data(),
-            dist.data(),
-            samples_per_batch,
-            n_exemplars,
-            n,
-            true);
-        break;
-      case raft::distance::DistanceType::L1:
-        raft::distance::distance<raft::distance::DistanceType::L1, value_t, value_t, value_t, int>(
-          handle,
-          query + batch_offset * n,
-          exemplars_dense.data(),
-          dist.data(),
-          samples_per_batch,
-          n_exemplars,
-          n,
-          true);
-        break;
-      case raft::distance::DistanceType::CosineExpanded:
-        raft::distance::
-          distance<raft::distance::DistanceType::CosineExpanded, value_t, value_t, value_t, int>(
-            handle,
-            query + batch_offset * n,
-            exemplars_dense.data(),
-            dist.data(),
-            samples_per_batch,
-            n_exemplars,
-            n,
-            true);
-        break;
-      default: RAFT_EXPECTS(false, "Incorrect metric passed!");
-    }
+    // compute the distances using the CUVS API
+    cuvs::distance::pairwise_distance(
+      handle,
+      raft::make_device_matrix_view<const value_t, int64_t>(
+        query + batch_offset * n, samples_per_batch, n),
+      raft::make_device_matrix_view<const value_t, int64_t>(exemplars_dense.data(), n_exemplars, n),
+      raft::make_device_matrix_view<value_t, int64_t>(dist.data(), samples_per_batch, n_exemplars),
+      static_cast<cuvs::distance::DistanceType>(metric));
 
     // compute the minimum distances to exemplars of each cluster
     value_idx n_elements = samples_per_batch * n_selected_clusters;