Benchmark CAGRA select_k

cpp/bench/prims/matrix/select_k.cu

@@ -273,6 +273,11 @@ const static size_t MAX_MEMORY = 16 * 1024 * 1024 * 1024ULL;
  SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kRadix8bits, input) \
  SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kRadix11bits, input) \
  SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kRadix11bitsExtraPass, input) \
+ if (input.k <= 1024) { \
+ if constexpr (std::is_same_v<KeyT, float>) { \
+ SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kCagra, input) \
+ } \
+ } \
  if (input.k <= raft::matrix::detail::select::warpsort::kMaxCapacity) { \
  SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kWarpImmediate, input) \
  SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kWarpFiltered, input) \

cpp/include/raft/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh

@@ -885,8 +885,8 @@ __launch_bounds__(1024, 1) RAFT_KERNEL
  if (i_batch >= size_batch) return;
 
  constexpr uint32_t smem_len = 2 * maxTopk + 2048 + 8;
- static_assert(maxTopk * (1 + utils::size_of<ValT>() / utils::size_of<uint32_t>()) <= smem_len,
- "maxTopk * sizeof(ValT) must be smaller or equal to 8192 byte");
+ // static_assert(maxTopk * (1 + utils::size_of<ValT>() / utils::size_of<uint32_t>()) <= smem_len,
+ // "maxTopk * sizeof(ValT) must be smaller or equal to 8192 byte");
  __shared__ uint32_t _smem[smem_len];
 
  topk_cta_11_core<stateBitLen, vecLen, maxTopk, numSortThreads, ValT>(

cpp/include/raft/neighbors/detail/cagra/utils.hpp

@@ -34,6 +34,11 @@ inline cudaDataType_t get_cuda_data_type<float>()
  return CUDA_R_32F;
 }
 template <>
+inline cudaDataType_t get_cuda_data_type<double>()
+{
+ return CUDA_R_64F;
+}
+template <>
 inline cudaDataType_t get_cuda_data_type<half>()
 {
  return CUDA_R_16F;
@@ -87,6 +92,11 @@ _RAFT_HOST_DEVICE constexpr unsigned size_of<std::uint64_t>()
  return 8;
 }
 template <>
+_RAFT_HOST_DEVICE constexpr unsigned size_of<std::int64_t>()
+{
+ return 8;
+}
+template <>
 _RAFT_HOST_DEVICE constexpr unsigned size_of<uint4>()
 {
  return 16;

cpp/internal/raft_internal/matrix/select_k.cuh

@@ -22,6 +22,9 @@
 #include <raft/matrix/detail/select_warpsort.cuh>
 #include <raft/matrix/select_k.cuh>
 #include <raft/neighbors/detail/selection_faiss.cuh>
+// #include <raft/neighbors/detail/cagra/topk_for_cagra/topk.h>
+#include <raft/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh>
+#include <raft/neighbors/detail/cagra/utils.hpp>
 
 namespace raft::matrix::select {
 
@@ -59,7 +62,8 @@ enum class Algo {
  kWarpFiltered,
  kWarpDistributed,
  kWarpDistributedShm,
- kFaissBlockSelect
+ kFaissBlockSelect,
+ kCagra
 };
 
 inline auto operator<<(std::ostream& os, const Algo& algo) -> std::ostream&
@@ -75,6 +79,7 @@ inline auto operator<<(std::ostream& os, const Algo& algo) -> std::ostream&
  case Algo::kWarpDistributed: return os << "kWarpDistributed";
  case Algo::kWarpDistributedShm: return os << "kWarpDistributedShm";
  case Algo::kFaissBlockSelect: return os << "kFaissBlockSelect";
+ case Algo::kCagra: return os << "kCagra";
  default: return os << "unknown enum value";
  }
 }
@@ -170,6 +175,33 @@ void select_k_impl(const resources& handle,
  case Algo::kFaissBlockSelect:
  return neighbors::detail::select_k(
  in, in_idx, batch_size, len, out, out_idx, select_min, k, stream);
+ case Algo::kCagra: {
+ // TODO: afaict cagra top-k only works on floats
+ if constexpr (std::is_same_v<T, float>) {
+ auto dtype = raft::neighbors::experimental::cagra::detail::utils::get_cuda_data_type<T>();
+ size_t buffer_size =
+ raft::neighbors::experimental::cagra::detail::_cuann_find_topk_bufferSize(
+ k, batch_size, len, dtype);
+ rmm::device_uvector<char> buffer(buffer_size, stream);
+ raft::neighbors::experimental::cagra::detail::_cuann_find_topk<IdxT>(k,
+ batch_size,
+ len,
+ in,
+ len,
+ in_idx,
+ len,
+ out,
+ len,
+ out_idx,
+ len,
+ buffer.data(),
+ false,
+ NULL,
+ stream);
+ } else {
+ throw std::logic_error("TODO: CAGRA topk for double/half");
+ }
+ }
  }
 }
 } // namespace raft::matrix::select