Integrate KNN implementation: ivf-pq

cpp/include/raft/spatial/knn/detail/ann_utils.cuh

@@ -151,6 +151,7 @@ inline void memzero(T* ptr, size_t n_elems, rmm::cuda_stream_view stream)
   }
 }
 
+namespace {
 __global__ void argmin_along_rows_kernel(uint32_t n_rows,
                                          uint32_t n_cols,
                                          const float* a,
@@ -185,6 +186,7 @@ __global__ void argmin_along_rows_kernel(uint32_t n_rows,
   }
   if (threadIdx.x == 0) { out[i] = shm_ids[0]; }
 }
+}  // namespace
 
 /**
  * @brief Find index of the smallest element in each row.
@@ -209,6 +211,7 @@ inline void argmin_along_rows(
   argmin_along_rows_kernel<<<n_rows, block_dim, 0, stream>>>(n_rows, n_cols, a, out);
 }
 
+namespace {
 __global__ void dots_along_rows_kernel(uint32_t n_rows, uint32_t n_cols, const float* a, float* out)
 {
   uint64_t i = threadIdx.y + (blockDim.y * blockIdx.x);
@@ -226,6 +229,7 @@ __global__ void dots_along_rows_kernel(uint32_t n_rows, uint32_t n_cols, const f
   sqsum += __shfl_xor_sync(0xffffffff, sqsum, 16);
   if (threadIdx.x == 0) { out[i] = sqsum; }
 }
+}  // namespace
 
 /**
  * @brief Square sum of values in each row (row-major matrix).
@@ -317,6 +321,7 @@ void accumulate_into_selected(size_t n_rows,
   }
 }
 
+namespace {
 __global__ void normalize_rows_kernel(uint32_t n_rows, uint32_t n_cols, float* a)
 {
   uint64_t i = threadIdx.y + (blockDim.y * blockIdx.x);
@@ -338,6 +343,7 @@ __global__ void normalize_rows_kernel(uint32_t n_rows, uint32_t n_cols, float* a
     a[j + n_cols * i] *= sqsum;
   }
 }
+}  // namespace
 
 /**
  * @brief Divide rows by their L2 norm (square root of sum of squares).
@@ -557,4 +563,58 @@ void copy_selected(uint64_t n_rows,
     default: RAFT_FAIL("All pointers must reside on the same side, host or device.");
   }
 }
+
+template <typename T, typename S>
+__global__ void copy_fill_kernel(uint32_t n_rows,
+                                 uint32_t n_cols,
+                                 const S* src,
+                                 uint32_t ld_src,
+                                 T* dst,
+                                 uint32_t ld_dst,
+                                 T fill_value)
+{
+  uint64_t gid = threadIdx.x + blockDim.x * blockIdx.x;
+  uint64_t j   = gid % ld_dst;
+  uint64_t i   = gid / ld_dst;
+  if (i >= n_rows) return;
+  dst[j + ld_dst * i] = j < n_cols ? utils::mapping<T>{}(src[j + ld_src * i]) : fill_value;
+}
+
+/**
+ * @brief Copy a number of consecutive rows filling the longer output row endings with the given
+ * constant.
+ *
+ * @tparam T target type
+ * @tparam S source type
+ *
+ * @param n_rows
+ * @param n_cols
+ * @param[in] src input matrix [n_rows, ld_src]
+ * @param ld_src number of cols in the input (ld_src >= n_cols)
+ * @param[out] dst output matrix [n_rows, ld_dst]
+ * @param ld_dst number of cols in the output (ld_dst >= n_cols)
+ * @param fill_value the value to fill at column indices [n_cols, ... ld_dst)
+ * @param stream
+ */
+template <typename T, typename S>
+void copy_fill(uint32_t n_rows,
+               uint32_t n_cols,
+               const S* src,
+               uint32_t ld_src,
+               T* dst,
+               uint32_t ld_dst,
+               T fill_value,
+               rmm::cuda_stream_view stream)
+{
+  RAFT_EXPECTS(ld_src >= n_cols, "src leading dimension must be larger than n_cols");
+  RAFT_EXPECTS(ld_dst >= n_cols, "dist leading dimension must be larger than n_cols");
+  dim3 threads(128, 1, 1);
+  dim3 blocks(
+    ceildiv<uint64_t>(static_cast<uint64_t>(n_rows) * static_cast<uint64_t>(ld_dst), threads.x),
+    1,
+    1);
+  copy_fill_kernel<T, S>
+    <<<blocks, threads, 0, stream>>>(n_rows, n_cols, src, ld_src, dst, ld_dst, fill_value);
+}
+
 }  // namespace raft::spatial::knn::detail::utils

cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh

@@ -18,7 +18,7 @@
 
 #include "../ivf_flat_types.hpp"
 #include "ann_utils.cuh"
-#include "topk/radix_topk.cuh"
+#include "topk.cuh"
 #include "topk/warpsort_topk.cuh"
 
 #include <raft/common/device_loads_stores.cuh>
@@ -1133,29 +1133,16 @@ void search_impl(const handle_t& handle,
                stream);
 
   RAFT_LOG_TRACE_VEC(distance_buffer_dev.data(), std::min<uint32_t>(20, index.n_lists()));
-  if (n_probes <= raft::spatial::knn::detail::topk::kMaxCapacity) {
-    topk::warp_sort_topk<AccT, uint32_t>(distance_buffer_dev.data(),
-                                         nullptr,
-                                         n_queries,
-                                         index.n_lists(),
-                                         n_probes,
-                                         coarse_distances_dev.data(),
-                                         coarse_indices_dev.data(),
-                                         select_min,
-                                         stream,
-                                         search_mr);
-  } else {
-    topk::radix_topk<AccT, uint32_t, 11, 512>(distance_buffer_dev.data(),
-                                              nullptr,
-                                              n_queries,
-                                              index.n_lists(),
-                                              n_probes,
-                                              coarse_distances_dev.data(),
-                                              coarse_indices_dev.data(),
-                                              select_min,
-                                              stream,
-                                              search_mr);
-  }
+  select_topk<AccT, uint32_t>(distance_buffer_dev.data(),
+                              nullptr,
+                              n_queries,
+                              index.n_lists(),
+                              n_probes,
+                              coarse_distances_dev.data(),
+                              coarse_indices_dev.data(),
+                              select_min,
+                              stream,
+                              search_mr);
   RAFT_LOG_TRACE_VEC(coarse_indices_dev.data(), n_probes);
   RAFT_LOG_TRACE_VEC(coarse_distances_dev.data(), n_probes);
 
@@ -1204,31 +1191,16 @@ void search_impl(const handle_t& handle,
 
   // Merge topk values from different blocks
   if (grid_dim_x > 1) {
-    if (k <= raft::spatial::knn::detail::topk::kMaxCapacity) {
-      topk::warp_sort_topk<AccT, IdxT>(refined_distances_dev.data(),
-                                       refined_indices_dev.data(),
-                                       n_queries,
-                                       k * grid_dim_x,
-                                       k,
-                                       distances,
-                                       neighbors,
-                                       select_min,
-                                       stream,
-                                       search_mr);
-    } else {
-      // NB: this branch can only be triggered once `ivfflat_interleaved_scan` above supports larger
-      // `k` values (kMaxCapacity limit as a dependency of topk::block_sort)
-      topk::radix_topk<AccT, IdxT, 11, 512>(refined_distances_dev.data(),
-                                            refined_indices_dev.data(),
-                                            n_queries,
-                                            k * grid_dim_x,
-                                            k,
-                                            distances,
-                                            neighbors,
-                                            select_min,
-                                            stream,
-                                            search_mr);
-    }
+    select_topk<AccT, IdxT>(refined_distances_dev.data(),
+                            refined_indices_dev.data(),
+                            n_queries,
+                            k * grid_dim_x,
+                            k,
+                            distances,
+                            neighbors,
+                            select_min,
+                            stream,
+                            search_mr);
   }
 }