Merge pull request #2122 from rapidsai/branch-24.02

Forward-merge branch-24.02 to branch-24.04

cpp/CMakeLists.txt

@@ -424,6 +424,7 @@ if(RAFT_COMPILE_LIBRARY)
     src/raft_runtime/neighbors/cagra_build.cu
     src/raft_runtime/neighbors/cagra_search.cu
     src/raft_runtime/neighbors/cagra_serialize.cu
+    src/raft_runtime/neighbors/eps_neighborhood.cu
     src/raft_runtime/neighbors/ivf_flat_build.cu
     src/raft_runtime/neighbors/ivf_flat_search.cu
     src/raft_runtime/neighbors/ivf_flat_serialize.cu
@@ -443,6 +444,7 @@ if(RAFT_COMPILE_LIBRARY)
     src/raft_runtime/random/rmat_rectangular_generator_int64_float.cu
     src/raft_runtime/random/rmat_rectangular_generator_int_double.cu
     src/raft_runtime/random/rmat_rectangular_generator_int_float.cu
+    src/spatial/knn/detail/ball_cover/registers_eps_pass_euclidean.cu
     src/spatial/knn/detail/ball_cover/registers_pass_one_2d_dist.cu
     src/spatial/knn/detail/ball_cover/registers_pass_one_2d_euclidean.cu
     src/spatial/knn/detail/ball_cover/registers_pass_one_2d_haversine.cu

cpp/include/raft/neighbors/ball_cover-ext.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-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.
@@ -67,6 +67,25 @@ void knn_query(raft::resources const& handle,
                bool perform_post_filtering = true,
                float weight                = 1.0) RAFT_EXPLICIT;
 
+template <typename idx_t, typename value_t, typename int_t, typename matrix_idx_t>
+void eps_nn(raft::resources const& handle,
+            const BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index,
+            raft::device_matrix_view<bool, matrix_idx_t, row_major> adj,
+            raft::device_vector_view<idx_t, matrix_idx_t> vd,
+            raft::device_matrix_view<const value_t, matrix_idx_t, row_major> query,
+            value_t eps) RAFT_EXPLICIT;
+
+template <typename idx_t, typename value_t, typename int_t, typename matrix_idx_t>
+void eps_nn(raft::resources const& handle,
+            const BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index,
+            raft::device_vector_view<idx_t, matrix_idx_t> adj_ia,
+            raft::device_vector_view<idx_t, matrix_idx_t> adj_ja,
+            raft::device_vector_view<idx_t, matrix_idx_t> vd,
+            raft::device_matrix_view<const value_t, matrix_idx_t, row_major> query,
+            value_t eps,
+            std::optional<raft::host_scalar_view<int_t, matrix_idx_t>> max_k = std::nullopt)
+  RAFT_EXPLICIT;
+
 }  // namespace raft::neighbors::ball_cover
 
 #endif  // RAFT_EXPLICIT_INSTANTIATE_ONLY
@@ -87,6 +106,24 @@ void knn_query(raft::resources const& handle,
     bool perform_post_filtering,                                                                   \
     float weight);                                                                                 \
                                                                                                    \
+  extern template void raft::neighbors::ball_cover::eps_nn<idx_t, value_t, int_t, matrix_idx_t>(   \
+    raft::resources const& handle,                                                                 \
+    const raft::neighbors::ball_cover::BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index, \
+    raft::device_matrix_view<bool, matrix_idx_t, row_major> adj,                                   \
+    raft::device_vector_view<idx_t, matrix_idx_t> vd,                                              \
+    raft::device_matrix_view<const value_t, matrix_idx_t, row_major> query,                        \
+    value_t eps);                                                                                  \
+                                                                                                   \
+  extern template void raft::neighbors::ball_cover::eps_nn<idx_t, value_t, int_t, matrix_idx_t>(   \
+    raft::resources const& handle,                                                                 \
+    const raft::neighbors::ball_cover::BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index, \
+    raft::device_vector_view<idx_t, matrix_idx_t> adj_ia,                                          \
+    raft::device_vector_view<idx_t, matrix_idx_t> adj_ja,                                          \
+    raft::device_vector_view<idx_t, matrix_idx_t> vd,                                              \
+    raft::device_matrix_view<const value_t, matrix_idx_t, row_major> query,                        \
+    value_t eps,                                                                                   \
+    std::optional<raft::host_scalar_view<int_t, matrix_idx_t>> max_k);                             \
+                                                                                                   \
   extern template void                                                                             \
   raft::neighbors::ball_cover::all_knn_query<idx_t, value_t, int_t, matrix_idx_t>(                 \
     raft::resources const& handle,                                                                 \
@@ -119,6 +156,6 @@ void knn_query(raft::resources const& handle,
     bool perform_post_filtering,                                                                   \
     float weight);
 
-instantiate_raft_neighbors_ball_cover(int64_t, float, uint32_t, uint32_t);
+instantiate_raft_neighbors_ball_cover(int64_t, float, int64_t, int64_t);
 
 #undef instantiate_raft_neighbors_ball_cover

cpp/include/raft/neighbors/ball_cover-inl.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-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.
@@ -63,7 +63,6 @@ template <typename idx_t, typename value_t, typename int_t, typename matrix_idx_
 void build_index(raft::resources const& handle,
                  BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index)
 {
-  ASSERT(index.n <= 3, "only 2d and 3d vectors are supported in current implementation");
   if (index.metric == raft::distance::DistanceType::Haversine) {
     raft::spatial::knn::detail::rbc_build_index(
       handle, index, spatial::knn::detail::HaversineFunc<value_t, int_t>());
@@ -255,9 +254,9 @@ void all_knn_query(raft::resources const& handle,
  *               looking in the closest landmark.
  * @param[in] n_query_pts number of query points
  */
-template <typename idx_t, typename value_t, typename int_t>
+template <typename idx_t, typename value_t, typename int_t, typename matrix_idx = std::int64_t>
 void knn_query(raft::resources const& handle,
-               const BallCoverIndex<idx_t, value_t, int_t>& index,
+               const BallCoverIndex<idx_t, value_t, int_t, matrix_idx>& index,
                int_t k,
                const value_t* query,
                int_t n_query_pts,
@@ -295,6 +294,106 @@ void knn_query(raft::resources const& handle,
   }
 }
 
+/**
+ * @brief Computes epsilon neighborhood for the L2 distance metric using rbc
+ *
+ * @tparam value_t   IO and math type
+ * @tparam idx_t    Index type
+ *
+ * @param[in] handle raft handle for resource management
+ * @param[in] index ball cover index which has been built
+ * @param[out] adj    adjacency matrix [row-major] [on device] [dim = m x n]
+ * @param[out] vd     vertex degree array [on device] [len = m + 1]
+ *                    `vd + m` stores the total number of edges in the adjacency
+ *                    matrix. Pass a nullptr if you don't need this info.
+ * @param[in]  query  first matrix [row-major] [on device] [dim = m x k]
+ * @param[in]  eps    defines epsilon neighborhood radius
+ */
+template <typename idx_t, typename value_t, typename int_t, typename matrix_idx_t = std::int64_t>
+void eps_nn(raft::resources const& handle,
+            const BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index,
+            raft::device_matrix_view<bool, matrix_idx_t, row_major> adj,
+            raft::device_vector_view<idx_t, matrix_idx_t> vd,
+            raft::device_matrix_view<const value_t, matrix_idx_t, row_major> query,
+            value_t eps)
+{
+  ASSERT(index.n == query.extent(1), "vector dimension needs to be the same for index and queries");
+  ASSERT(index.metric == raft::distance::DistanceType::L2SqrtExpanded ||
+           index.metric == raft::distance::DistanceType::L2SqrtUnexpanded,
+         "Metric not supported");
+  ASSERT(index.is_index_trained(), "index must be previously trained");
+
+  // run query
+  raft::spatial::knn::detail::rbc_eps_nn_query(
+    handle,
+    index,
+    eps,
+    query.data_handle(),
+    query.extent(0),
+    adj.data_handle(),
+    vd.data_handle(),
+    spatial::knn::detail::EuclideanFunc<value_t, int_t>());
+}
+
+/**
+ * @brief Computes epsilon neighborhood for the L2 distance metric using rbc
+ *
+ * @tparam value_t   IO and math type
+ * @tparam idx_t    Index type
+ *
+ * @param[in] handle raft handle for resource management
+ * @param[in] index ball cover index which has been built
+ * @param[out] adj_ia    adjacency matrix CSR row offsets
+ * @param[out] adj_ja    adjacency matrix CSR column indices, needs to be nullptr
+ *                       in first pass with max_k nullopt
+ * @param[out] vd     vertex degree array [on device] [len = m + 1]
+ *                    `vd + m` stores the total number of edges in the adjacency
+ *                    matrix. Pass a nullptr if you don't need this info.
+ * @param[in]  query  first matrix [row-major] [on device] [dim = m x k]
+ * @param[in]  eps    defines epsilon neighborhood radius
+ * @param[inout] max_k if nullopt (default), the user needs to make 2 subsequent calls:
+ *                     The first call computes row offsets in adj_ia, where adj_ia[m]
+ *                     contains the minimum required size for adj_ja.
+ *                     The second call fills in adj_ja based on adj_ia.
+ *                     If max_k != nullopt the algorithm only fills up neighbors up to a
+ *                     maximum number of max_k for each row in a single pass. Note
+ *                     that it is not guarantueed to return the nearest neighbors.
+ *                     Upon return max_k is overwritten with the actual max_k found during
+ *                     computation.
+ */
+template <typename idx_t, typename value_t, typename int_t, typename matrix_idx_t = std::int64_t>
+void eps_nn(raft::resources const& handle,
+            const BallCoverIndex<idx_t, value_t, int_t, matrix_idx_t>& index,
+            raft::device_vector_view<idx_t, matrix_idx_t> adj_ia,
+            raft::device_vector_view<idx_t, matrix_idx_t> adj_ja,
+            raft::device_vector_view<idx_t, matrix_idx_t> vd,
+            raft::device_matrix_view<const value_t, matrix_idx_t, row_major> query,
+            value_t eps,
+            std::optional<raft::host_scalar_view<int_t, matrix_idx_t>> max_k = std::nullopt)
+{
+  ASSERT(index.n == query.extent(1), "vector dimension needs to be the same for index and queries");
+  ASSERT(index.metric == raft::distance::DistanceType::L2SqrtExpanded ||
+           index.metric == raft::distance::DistanceType::L2SqrtUnexpanded,
+         "Metric not supported");
+  ASSERT(index.is_index_trained(), "index must be previously trained");
+
+  int_t* max_k_ptr = nullptr;
+  if (max_k.has_value()) { max_k_ptr = max_k.value().data_handle(); }
+
+  // run query
+  raft::spatial::knn::detail::rbc_eps_nn_query(
+    handle,
+    index,
+    eps,
+    max_k_ptr,
+    query.data_handle(),
+    query.extent(0),
+    adj_ia.data_handle(),
+    adj_ja.data_handle(),
+    vd.data_handle(),
+    spatial::knn::detail::EuclideanFunc<value_t, int_t>());
+}
+
 /**
  * @ingroup random_ball_cover
  * @{
@@ -377,7 +476,7 @@ void knn_query(raft::resources const& handle,
             index,
             k,
             query.data_handle(),
-            query.extent(0),
+            (int_t)query.extent(0),
             inds.data_handle(),
             dists.data_handle(),
             perform_post_filtering,

cpp/include/raft/neighbors/ball_cover.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-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.
@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 #pragma once
-
 #ifndef RAFT_EXPLICIT_INSTANTIATE_ONLY
 #include "ball_cover-inl.cuh"
 #endif

cpp/include/raft/neighbors/ball_cover_types.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2021-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.
@@ -19,6 +19,7 @@
 #include <cstdint>
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_mdspan.hpp>
+#include <raft/core/host_mdspan.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/distance/distance_types.hpp>
 #include <rmm/device_uvector.hpp>
@@ -41,8 +42,8 @@ namespace raft::neighbors::ball_cover {
  */
 template <typename value_idx,
           typename value_t,
-          typename value_int  = std::uint32_t,
-          typename matrix_idx = std::uint32_t>
+          typename value_int  = std::int64_t,
+          typename matrix_idx = std::int64_t>
 class BallCoverIndex {
  public:
   explicit BallCoverIndex(raft::resources const& handle_,

cpp/include/raft/neighbors/brute_force-ext.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2023, NVIDIA CORPORATION.
+ * Copyright (c) 2020-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.

cpp/include/raft/neighbors/detail/knn_brute_force.cuh

@@ -443,13 +443,13 @@ void brute_force_knn_impl(
       if (metric == raft::distance::DistanceType::L2SqrtExpanded ||
           metric == raft::distance::DistanceType::L2SqrtUnexpanded ||
           metric == raft::distance::DistanceType::LpUnexpanded) {
-        float p = 0.5;  // standard l2
+        value_t p = 0.5;  // standard l2
         if (metric == raft::distance::DistanceType::LpUnexpanded) p = 1.0 / metricArg;
-        raft::linalg::unaryOp<float>(
+        raft::linalg::unaryOp<value_t>(
           res_D,
           res_D,
           n * k,
-          [p] __device__(float input) { return powf(fabsf(input), p); },
+          [p] __device__(value_t input) { return powf(fabsf(input), p); },
           stream);
       }
     } else {

cpp/include/raft/neighbors/detail/knn_merge_parts.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-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.
@@ -111,7 +111,7 @@ inline void knn_merge_parts_impl(const value_t* inK,
 {
   auto grid = dim3(n_samples);
 
-  constexpr int n_threads = (warp_q <= 1024) ? 128 : 64;
+  constexpr int n_threads = (warp_q < 1024) ? 128 : 64;
   auto block              = dim3(n_threads);
 
   auto kInit = std::numeric_limits<value_t>::max();