[WIP] Share RAFT selection kernel instantiations

cpp/CMakeLists.txt

@@ -1441,6 +1441,66 @@ if(NOT BUILD_CPU_ONLY)
     "\n"
   )
 
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_dense_launch_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_dense_launch.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_dense_launch_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE
+      "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_dense_view_launch_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_dense_view_launch.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_dense_view_launch_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_radix_launch_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_radix_launch.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_radix_launch_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE
+      "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_segmented_sort_launch_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_segmented_sort_launch.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_segmented_sort_launch_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE
+      "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_warpsort_launch_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_warpsort_launch.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_warpsort_launch_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE
+      "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_warpsort_parameters_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_warpsort_parameters.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_warpsort_parameters_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_warpsort_setup_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/selection/select_k_warpsort_setup.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_warpsort_setup_@instantiation_name@.cu"
+  )
+
   generate_inst_matrix(
     cuvs_primitive_sources
     MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/add_matrix.json"
@@ -1502,6 +1562,7 @@ if(NOT BUILD_CPU_ONLY)
     cuvs_objs OBJECT
     ${cuvs_primitive_sources}
     src/core/cub_ptx_query.cu
+    src/selection/select_k_csr_launch.cu
     src/cluster/detail/minClusterDistanceCompute.cu
     src/cluster/agglomerative.cu
     ${gmm_inst_files}

cpp/src/core/selection_launchers.hpp

@@ -0,0 +1,78 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cuda/stream>
+#include <optional>
+#include <raft/core/device_csr_matrix.hpp>
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/resources.hpp>
+#include <raft/matrix/select_k_types.hpp>
+
+namespace cuvs::detail {
+template <typename T, typename Index>
+void select_k(raft::resources const& handle,
+              T const* values,
+              Index const* indices,
+              size_t rows,
+              size_t columns,
+              int k,
+              T* selected_values,
+              Index* selected_indices,
+              bool select_min,
+              bool sorted,
+              raft::matrix::SelectAlgo algorithm,
+              Index const* row_lengths);
+
+template <typename T, typename Index>
+void select_k(
+  raft::resources const& handle,
+  raft::device_matrix_view<const T, int64_t, raft::row_major> values,
+  std::optional<raft::device_matrix_view<const Index, int64_t, raft::row_major>> indices,
+  raft::device_matrix_view<T, int64_t, raft::row_major> selected_values,
+  raft::device_matrix_view<Index, int64_t, raft::row_major> selected_indices,
+  bool select_min,
+  bool sorted                        = false,
+  raft::matrix::SelectAlgo algorithm = raft::matrix::SelectAlgo::kAuto);
+
+template <typename T, typename Index>
+void select_k(raft::resources const& handle,
+              raft::device_csr_matrix_view<const T, Index, Index, Index> values,
+              std::optional<raft::device_vector_view<const Index, Index>> indices,
+              raft::device_matrix_view<T, Index, raft::row_major> selected_values,
+              raft::device_matrix_view<Index, Index, raft::row_major> selected_indices,
+              bool select_min,
+              bool sorted                        = false,
+              raft::matrix::SelectAlgo algorithm = raft::matrix::SelectAlgo::kAuto);
+
+struct warpsort_launch_params {
+  int block_size;
+  int min_grid_size;
+};
+
+template <template <int, bool, typename, typename> class WarpSortClass,
+          typename T,
+          typename IdxT,
+          typename RowLayout>
+struct warpsort_launch_setup {
+  static warpsort_launch_params calc_optimal_params(int k, int block_size_limit);
+
+  static void kernel(int k,
+                     bool select_min,
+                     size_t batch_size,
+                     size_t len,
+                     int num_blocks,
+                     int block_dim,
+                     size_t smem_size,
+                     const T* in_key,
+                     const IdxT* in_idx,
+                     const IdxT* in_indptr,
+                     T* out_key,
+                     IdxT* out_idx,
+                     cuda::stream_ref stream);
+};
+
+}  // namespace cuvs::detail

cpp/src/neighbors/detail/cagra/cagra_search.cuh

@@ -7,6 +7,7 @@
 
 #include "../../../core/map_launchers.hpp"
 #include "../../../core/nvtx.hpp"
+#include "../../../core/selection_launchers.hpp"
 #include "factory.cuh"
 #include "sample_filter_utils.cuh"
 #include "search_multi_cta.cuh"
@@ -36,7 +37,6 @@
 #include <raft/linalg/matrix_vector_op.cuh>
 #include <raft/linalg/norm.cuh>
 #include <raft/linalg/reduce.cuh>
-#include <raft/matrix/select_k.cuh>
 
 #include <cstddef>
 
@@ -599,7 +599,7 @@ void search_multi_partition(
     // Cosine transform applied here) are smaller-is-closer, while InnerProduct is larger-is-closer.
     const bool select_min = cuvs::distance::is_min_close(metric);
 
-    raft::matrix::select_k<DistanceT, uint32_t>(
+    cuvs::detail::select_k<DistanceT, uint32_t>(
       res,
       raft::make_device_matrix_view<const DistanceT, int64_t, raft::row_major>(
         transposed_distances,

cpp/src/neighbors/detail/knn_brute_force.cuh

@@ -5,6 +5,7 @@
 
 #pragma once
 
+#include "../../core/selection_launchers.hpp"
 #include <cuvs/distance/distance.hpp>
 #include <cuvs/neighbors/brute_force.hpp>
 #include <cuvs/neighbors/knn_merge_parts.hpp>
@@ -33,7 +34,6 @@
 #include <raft/sparse/convert/coo.cuh>
 #include <raft/sparse/convert/csr.cuh>
 #include <raft/sparse/linalg/masked_matmul.cuh>
-#include <raft/sparse/matrix/select_k.cuh>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 #include <raft/util/popc.cuh>
@@ -897,8 +897,7 @@ void brute_force_search_filtered(
       csr.get_elements().data(), compressed_csr_view);
     std::optional<raft::device_vector_view<const IdxT, IdxT>> no_opt = std::nullopt;
     bool select_min = cuvs::distance::is_min_close(metric);
-    raft::sparse::matrix::select_k(
-      res, const_csr_view, no_opt, distances, neighbors, select_min, true);
+    cuvs::detail::select_k(res, const_csr_view, no_opt, distances, neighbors, select_min, true);
   }
 
   return;

cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh

@@ -7,6 +7,7 @@
 
 #include "../../core/map_launchers.hpp"
 #include "../../core/nvtx.hpp"
+#include "../../core/selection_launchers.hpp"
 #include "../detail/ann_utils.cuh"
 #include "../ivf_common.cuh"
 #include "../sample_filter.cuh"  // none_sample_filter
@@ -30,7 +31,6 @@
 #include <raft/linalg/norm_types.hpp>
 #include <raft/linalg/normalize.cuh>
 #include <raft/matrix/detail/select_warpsort.cuh>
-#include <raft/matrix/select_k.cuh>
 #include <raft/util/cache.hpp>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
@@ -245,7 +245,7 @@ void select_clusters(raft::resources const& handle,
   // Select neighbor clusters for each query.
   rmm::device_uvector<dist_type> cluster_dists(size_t(n_queries) * size_t(n_probes), stream, mr);
   // cuvs::selection::select_k lacks uint32_t-as-a-value support at the moment
-  raft::matrix::select_k<dist_type, uint32_t>(
+  cuvs::detail::select_k<dist_type, uint32_t>(
     handle,
     raft::make_device_matrix_view<const dist_type, int64_t>(
       qc_distances.data(), n_queries, n_lists),

cpp/src/selection/select_k.cuh

@@ -5,8 +5,8 @@
 
 #pragma once
 
+#include "../core/selection_launchers.hpp"
 #include <cuvs/selection/select_k.hpp>
-#include <raft/matrix/detail/select_k.cuh>
 
 namespace cuvs::selection::detail {
 
@@ -36,19 +36,18 @@ void select_k(raft::resources const& handle,
 
   // just delegate implementation to raft - the primary benefit here is to have
   // instantiations only compiled once in cuvs
-  return raft::matrix::detail::select_k<T, IdxT>(
-    handle,
-    in_val.data_handle(),
-    in_idx.has_value() ? in_idx->data_handle() : nullptr,
-    batch_size,
-    len,
-    k,
-    out_val.data_handle(),
-    out_idx.data_handle(),
-    select_min,
-    sorted,
-    algo,
-    len_i.has_value() ? len_i->data_handle() : nullptr);
+  return cuvs::detail::select_k<T, IdxT>(handle,
+                                         in_val.data_handle(),
+                                         in_idx.has_value() ? in_idx->data_handle() : nullptr,
+                                         batch_size,
+                                         len,
+                                         k,
+                                         out_val.data_handle(),
+                                         out_idx.data_handle(),
+                                         select_min,
+                                         sorted,
+                                         algo,
+                                         len_i.has_value() ? len_i->data_handle() : nullptr);
 }
 }  // namespace cuvs::selection::detail