From 0e9f61c765a557272f34845d5c38e0f2a5fb3de0 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Tue, 4 Nov 2025 00:30:20 -0800 Subject: [PATCH 01/25] covnert max_topk to a runtime parameter --- cpp/src/neighbors/detail/cagra/bitonic.hpp | 198 +++++-------- .../cagra/search_multi_cta_kernel-inl.cuh | 2 +- .../detail/cagra/search_single_cta.cuh | 8 +- .../cagra/search_single_cta_kernel-inl.cuh | 268 +++++++++--------- .../neighbors/detail/cagra/topk_by_radix.cuh | 73 ++--- .../detail/cagra/topk_for_cagra/topk.cu | 124 ++++---- .../detail/cagra/topk_for_cagra/topk_core.cuh | 103 +++---- 7 files changed, 350 insertions(+), 426 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/bitonic.hpp b/cpp/src/neighbors/detail/cagra/bitonic.hpp index 66c726c71a..6a003face9 100644 --- a/cpp/src/neighbors/detail/cagra/bitonic.hpp +++ b/cpp/src/neighbors/detail/cagra/bitonic.hpp @@ -18,12 +18,8 @@ template RAFT_DEVICE_INLINE_FUNCTION void swap_if_needed(K& k0, V& v0, K& k1, V& v1, const bool asc) { if ((k0 != k1) && ((k0 < k1) != asc)) { - const auto tmp_k = k0; - k0 = k1; - k1 = tmp_k; - const auto tmp_v = v0; - v0 = v1; - v1 = tmp_v; + cuda::std::swap(k0, k1); + cuda::std::swap(v0, v1); } } @@ -41,106 +37,52 @@ RAFT_DEVICE_INLINE_FUNCTION void swap_if_needed(K& k0, } } -template -struct warp_merge_core { - RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[N], - V v[N], - const std::uint32_t range, - const bool asc) +template +struct warp_merge_core_n { + RAFT_DEVICE_INLINE_FUNCTION void operator()( + K* ks, V* vs, unsigned n, const std::uint32_t range, const bool asc) { const auto lane_id = threadIdx.x % warp_size; if (range == 1) { - for (std::uint32_t b = 2; b <= N; b <<= 1) { + for (std::uint32_t b = 2; b <= n; b <<= 1) { for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { #pragma unroll - for (std::uint32_t i = 0; i < N; i++) { + for (std::uint32_t i = 0; i < n; i++) { std::uint32_t j = i ^ c; if (i >= j) continue; - const auto line_id = i + (N * lane_id); + const auto line_id = i + (n * lane_id); const auto p = static_cast(line_id & b) == static_cast(line_id & c); - swap_if_needed(k[i], v[i], k[j], v[j], p); + swap_if_needed(ks[i], vs[i], ks[j], vs[j], p); } } } - return; - } - - const std::uint32_t b = range; - for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { - const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); -#pragma unroll - for (std::uint32_t i = 0; i < N; i++) { - swap_if_needed(k[i], v[i], c, p); - } - } - const auto p = ((lane_id & b) == 0); - for (std::uint32_t c = N / 2; c >= 1; c >>= 1) { + } else { + const std::uint32_t b = range; + for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { + const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); #pragma unroll - for (std::uint32_t i = 0; i < N; i++) { - std::uint32_t j = i ^ c; - if (i >= j) continue; - swap_if_needed(k[i], v[i], k[j], v[j], p); - } - } - } -}; - -template -struct warp_merge_core { - RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[6], - V v[6], - const std::uint32_t range, - const bool asc) - { - constexpr unsigned N = 6; - const auto lane_id = threadIdx.x % warp_size; - - if (range == 1) { - for (std::uint32_t i = 0; i < N; i += 3) { - const auto p = (i == 0); - swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); - swap_if_needed(k[1 + i], v[1 + i], k[2 + i], v[2 + i], p); - swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); - } - const auto p = ((lane_id & 1) == 0); - for (std::uint32_t i = 0; i < 3; i++) { - std::uint32_t j = i + 3; - swap_if_needed(k[i], v[i], k[j], v[j], p); - } - for (std::uint32_t i = 0; i < N; i += 3) { - swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); - swap_if_needed(k[1 + i], v[1 + i], k[2 + i], v[2 + i], p); - swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); + for (std::uint32_t i = 0; i < n; i++) { + swap_if_needed(ks[i], vs[i], c, p); + } } - return; - } - - const std::uint32_t b = range; - for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { - const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); + const auto p = ((lane_id & b) == 0); + for (std::uint32_t c = n / 2; c >= 1; c >>= 1) { #pragma unroll - for (std::uint32_t i = 0; i < N; i++) { - swap_if_needed(k[i], v[i], c, p); + for (std::uint32_t i = 0; i < n; i++) { + std::uint32_t j = i ^ c; + if (i >= j) continue; + swap_if_needed(ks[i], vs[i], ks[j], vs[j], p); + } } } - const auto p = ((lane_id & b) == 0); - for (std::uint32_t i = 0; i < 3; i++) { - std::uint32_t j = i + 3; - swap_if_needed(k[i], v[i], k[j], v[j], p); - } - for (std::uint32_t i = 0; i < N; i += N / 2) { - swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); - swap_if_needed(k[1 + i], v[1 + i], k[2 + i], v[2 + i], p); - swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); - } } }; template -struct warp_merge_core { - RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[3], - V v[3], +struct warp_merge_core_3 { + RAFT_DEVICE_INLINE_FUNCTION void operator()(K* ks, + V* vs, const std::uint32_t range, const bool asc) { @@ -149,31 +91,30 @@ struct warp_merge_core { if (range == 1) { const auto p = ((lane_id & 1) == 0); - swap_if_needed(k[0], v[0], k[1], v[1], p); - swap_if_needed(k[1], v[1], k[2], v[2], p); - swap_if_needed(k[0], v[0], k[1], v[1], p); - return; - } - - const std::uint32_t b = range; - for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { - const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); + swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); + swap_if_needed(ks[1], vs[1], ks[2], vs[2], p); + swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); + } else { + const std::uint32_t b = range; + for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { + const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); #pragma unroll - for (std::uint32_t i = 0; i < N; i++) { - swap_if_needed(k[i], v[i], c, p); + for (std::uint32_t i = 0; i < N; i++) { + swap_if_needed(ks[i], vs[i], c, p); + } } + const auto p = ((lane_id & b) == 0); + swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); + swap_if_needed(ks[1], vs[1], ks[2], vs[2], p); + swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); } - const auto p = ((lane_id & b) == 0); - swap_if_needed(k[0], v[0], k[1], v[1], p); - swap_if_needed(k[1], v[1], k[2], v[2], p); - swap_if_needed(k[0], v[0], k[1], v[1], p); } }; template -struct warp_merge_core { - RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[2], - V v[2], +struct warp_merge_core_2 { + RAFT_DEVICE_INLINE_FUNCTION void operator()(K* ks, + V* vs, const std::uint32_t range, const bool asc) { @@ -182,27 +123,26 @@ struct warp_merge_core { if (range == 1) { const auto p = ((lane_id & 1) == 0); - swap_if_needed(k[0], v[0], k[1], v[1], p); - return; - } - - const std::uint32_t b = range; - for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { - const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); + swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); + } else { + const std::uint32_t b = range; + for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { + const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); #pragma unroll - for (std::uint32_t i = 0; i < N; i++) { - swap_if_needed(k[i], v[i], c, p); + for (std::uint32_t i = 0; i < N; i++) { + swap_if_needed(ks[i], vs[i], c, p); + } } + const auto p = ((lane_id & b) == 0); + swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); } - const auto p = ((lane_id & b) == 0); - swap_if_needed(k[0], v[0], k[1], v[1], p); } }; template -struct warp_merge_core { - RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[1], - V v[1], +struct warp_merge_core_1 { + RAFT_DEVICE_INLINE_FUNCTION void operator()(K* ks, + V* vs, const std::uint32_t range, const bool asc) { @@ -210,25 +150,33 @@ struct warp_merge_core { const std::uint32_t b = range; for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); - swap_if_needed(k[0], v[0], c, p); + swap_if_needed(ks[0], vs[0], c, p); } } }; } // namespace detail -template -RAFT_DEVICE_INLINE_FUNCTION void warp_merge(K k[N], V v[N], unsigned range, const bool asc = true) +template +RAFT_DEVICE_INLINE_FUNCTION void warp_merge( + K* ks, V* vs, unsigned n, unsigned range, const bool asc = true) { - detail::warp_merge_core{}(k, v, range, asc); + if (n == 1) { + detail::warp_merge_core_1{}(ks, vs, range, asc); + } else if (n == 2) { + detail::warp_merge_core_2{}(ks, vs, range, asc); + } else if (n == 3) { + detail::warp_merge_core_3{}(ks, vs, range, asc); + } else { + detail::warp_merge_core_n{}(ks, vs, n, range, asc); + } } -template -RAFT_DEVICE_INLINE_FUNCTION void warp_sort(K k[N], V v[N], const bool asc = true) +template +RAFT_DEVICE_INLINE_FUNCTION void warp_sort(K* ks, V* vs, unsigned n, const bool asc = true) { -#pragma unroll for (std::uint32_t range = 1; range <= warp_size; range <<= 1) { - warp_merge(k, v, range, asc); + warp_merge(ks, vs, n, range, asc); } } diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 1720841c91..7a50edc884 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -131,7 +131,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort(float* distances, // [num } } /* Warp Sort */ - bitonic::warp_sort(key, val); + bitonic::warp_sort(key, val, N); /* Store sorted results */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * lane_id) + i; diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh index 948c1b889e..9d59a74df2 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh @@ -133,9 +133,9 @@ struct search // Tentatively calculate the required share memory size when radix // sort based topk is used, assuming the block size is the maximum. if (itopk_size <= 256) { - additional_smem_size += topk_by_radix_sort<256, INDEX_T>::smem_size * sizeof(std::uint32_t); + additional_smem_size += topk_by_radix_sort::smem_size(256) * sizeof(std::uint32_t); } else { - additional_smem_size += topk_by_radix_sort<512, INDEX_T>::smem_size * sizeof(std::uint32_t); + additional_smem_size += topk_by_radix_sort::smem_size(512) * sizeof(std::uint32_t); } } @@ -197,10 +197,10 @@ struct search smem_size = base_smem_size; if (itopk_size <= 256) { constexpr unsigned MAX_ITOPK = 256; - smem_size += topk_by_radix_sort::smem_size * sizeof(std::uint32_t); + smem_size += topk_by_radix_sort::smem_size(MAX_ITOPK) * sizeof(std::uint32_t); } else { constexpr unsigned MAX_ITOPK = 512; - smem_size += topk_by_radix_sort::smem_size * sizeof(std::uint32_t); + smem_size += topk_by_radix_sort::smem_size(MAX_ITOPK) * sizeof(std::uint32_t); } } RAFT_LOG_DEBUG("# smem_size: %u", smem_size); diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 939589c942..7e2cbc55bb 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -107,16 +107,16 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( const std::uint32_t num_itopk, unsigned MULTI_WARPS = 0) { - const unsigned lane_id = threadIdx.x % 32; - const unsigned warp_id = threadIdx.x / 32; + const unsigned lane_id = threadIdx.x % raft::warp_size(); + const unsigned warp_id = threadIdx.x / raft::warp_size(); if (MULTI_WARPS == 0) { if (warp_id > 0) { return; } - constexpr unsigned N = (MAX_CANDIDATES + 31) / 32; + constexpr unsigned N = (MAX_CANDIDATES + (raft::warp_size() - 1)) / raft::warp_size(); float key[N]; IdxT val[N]; /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { - unsigned j = lane_id + (32 * i); + unsigned j = lane_id + (raft::warp_size() * i); if (j < num_candidates) { key[i] = candidate_distances[j]; val[i] = candidate_indices[j]; @@ -126,7 +126,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } /* Sort */ - bitonic::warp_sort(key, val); + bitonic::warp_sort(key, val, N); /* Reg -> Temp_itopk */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * lane_id) + i; @@ -138,13 +138,13 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } else { // Use two warps (64 threads) constexpr unsigned max_candidates_per_warp = (MAX_CANDIDATES + 1) / 2; - constexpr unsigned N = (max_candidates_per_warp + 31) / 32; + constexpr unsigned N = (max_candidates_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); float key[N]; IdxT val[N]; if (warp_id < 2) { /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { - unsigned jl = lane_id + (32 * i); + unsigned jl = lane_id + (raft::warp_size() * i); unsigned j = jl + (max_candidates_per_warp * warp_id); if (j < num_candidates) { key[i] = candidate_distances[j]; @@ -155,7 +155,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } /* Sort */ - bitonic::warp_sort(key, val); + bitonic::warp_sort(key, val, N); /* Reg -> Temp_candidates */ for (unsigned i = 0; i < N; i++) { unsigned jl = (N * lane_id) + i; @@ -188,7 +188,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( if (num_warps_used > 1) { __syncthreads(); } if (warp_id < num_warps_used) { /* Merge */ - bitonic::warp_merge(key, val, 32); + bitonic::warp_merge(key, val, N, raft::warp_size()); /* Reg -> Temp_itopk */ for (unsigned i = 0; i < N; i++) { unsigned jl = (N * lane_id) + i; @@ -203,10 +203,11 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] + const std::uint32_t max_itopk, const std::uint32_t num_itopk, float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] @@ -215,17 +216,23 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( const bool first, unsigned MULTI_WARPS = 0) { - const unsigned lane_id = threadIdx.x % 32; - const unsigned warp_id = threadIdx.x / 32; + const unsigned lane_id = threadIdx.x % raft::warp_size(); + const unsigned warp_id = threadIdx.x / raft::warp_size(); + + constexpr unsigned MAX_N = + 32; // if MAX_N >> N, we may have occupancy reduction as this inflates per-thread stack size. + // If this has significant performance impact, we may get memory space for key & val from + // dynamically sized shared memory. + float key[MAX_N]; + IdxT val[MAX_N]; if (MULTI_WARPS == 0) { if (warp_id > 0) { return; } - constexpr unsigned N = (MAX_ITOPK + 31) / 32; - float key[N]; - IdxT val[N]; + const unsigned N = (max_itopk + (raft::warp_size() - 1)) / raft::warp_size(); + assert(N <= MAX_N); if (first) { /* Load itopk results */ for (unsigned i = 0; i < N; i++) { - unsigned j = lane_id + (32 * i); + unsigned j = lane_id + (raft::warp_size() * i); if (j < num_itopk) { key[i] = itopk_distances[j]; val[i] = itopk_indices[j]; @@ -235,7 +242,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Sort */ - bitonic::warp_sort(key, val); + bitonic::warp_sort(key, val, N); } else { /* Load itopk results */ for (unsigned i = 0; i < N; i++) { @@ -251,8 +258,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } /* Merge candidates */ for (unsigned i = 0; i < N; i++) { - unsigned j = (N * lane_id) + i; // [0:MAX_ITOPK-1] - unsigned k = MAX_ITOPK - 1 - j; + unsigned j = (N * lane_id) + i; // [0:max_itopk-1] + unsigned k = max_itopk - 1 - j; if (k >= num_itopk || k >= num_candidates) continue; float candidate_key = candidate_distances[device::swizzling(k)]; if (key[i] > candidate_key) { @@ -261,7 +268,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Merge */ - bitonic::warp_merge(key, val, 32); + bitonic::warp_merge(key, val, N, raft::warp_size()); /* Store new itopk results */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * lane_id) + i; @@ -272,15 +279,14 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } else { // Use two warps (64 threads) or more - constexpr unsigned max_itopk_per_warp = (MAX_ITOPK + 1) / 2; - constexpr unsigned N = (max_itopk_per_warp + 31) / 32; - float key[N]; - IdxT val[N]; + const unsigned max_itopk_per_warp = (max_itopk + 1) / 2; + const unsigned N = (max_itopk_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); + assert(N <= MAX_N); if (first) { /* Load itop results (not sorted) */ if (warp_id < 2) { for (unsigned i = 0; i < N; i++) { - unsigned j = lane_id + (32 * i) + (max_itopk_per_warp * warp_id); + unsigned j = lane_id + (raft::warp_size() * i) + (max_itopk_per_warp * warp_id); if (j < num_itopk) { key[i] = itopk_distances[j]; val[i] = itopk_indices[j]; @@ -290,7 +296,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Sort */ - bitonic::warp_sort(key, val); + bitonic::warp_sort(key, val, N); /* Store intermedidate results */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * threadIdx.x) + i; @@ -304,7 +310,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( /* Load intermedidate results */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * threadIdx.x) + i; - unsigned k = MAX_ITOPK - 1 - j; + unsigned k = max_itopk - 1 - j; if (k >= num_itopk) continue; float temp_key = itopk_distances[device::swizzling(k)]; if (key[i] == temp_key) continue; @@ -314,7 +320,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Merge */ - bitonic::warp_merge(key, val, 32); + bitonic::warp_merge(key, val, N, raft::warp_size()); } __syncthreads(); /* Store itopk results (sorted) */ @@ -383,8 +389,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( unsigned j = (N * lane_id) + i; if (j < turning_point) { k = j + (num_itopk_div2 * warp_id); - } else if (j >= (MAX_ITOPK / 2 - num_itopk_div2)) { - j -= (MAX_ITOPK / 2 - num_itopk_div2); + } else if (j >= (max_itopk / 2 - num_itopk_div2)) { + j -= (max_itopk / 2 - num_itopk_div2); if ((turning_point <= j) && (j < num_itopk_div2)) { k = j + (num_itopk_div2 * warp_id); } } if (k < num_itopk) { @@ -396,7 +402,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Merge */ - bitonic::warp_merge(key, val, 32); + bitonic::warp_merge(key, val, N, raft::warp_size()); /* Store new itopk results */ for (unsigned i = 0; i < N; i++) { const unsigned j = (N * lane_id) + i; @@ -410,12 +416,12 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } -template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] + const std::uint32_t max_itopk, const std::uint32_t num_itopk, float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] @@ -431,15 +437,16 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( // The results sorted above are merged with the internal intermediate top-k // results so far using bitonic merge. - topk_by_bitonic_sort_and_merge(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - MULTI_WARPS_2); + topk_by_bitonic_sort_and_merge(itopk_distances, + itopk_indices, + max_itopk, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + MULTI_WARPS_2); } // This function move the invalid index element to the end of the itopk list. @@ -502,7 +509,6 @@ RAFT_DEVICE_INLINE_FUNCTION void hashmap_restore(INDEX_T* const hashmap_ptr, /** * @brief Search operation for a single query using a single thread block. * * - * @tparam MAX_ITOPK Maximum for the internal_topk argument. * @tparam MAX_CANDIDATES * @tparam TOPK_BY_BITONIC_SORT * @tparam DATASET_DESCRIPTOR_T @@ -533,8 +539,7 @@ RAFT_DEVICE_INLINE_FUNCTION void hashmap_restore(INDEX_T* const hashmap_ptr, * @param small_hash_reset_interval Interval for resetting the small hash. * @param query_id sequential id of the query in the batch */ -template = 64) && (MAX_CANDIDATES > 128)) ? 1 : 0; - const unsigned multi_warps_2 = ((blockDim.x >= 64) && (MAX_ITOPK > 256)) ? 1 : 0; + const unsigned multi_warps_2 = ((blockDim.x >= 64) && (max_itopk > 256)) ? 1 : 0; // reset small-hash table. if ((iter + 1) % small_hash_reset_interval == 0) { @@ -709,34 +715,35 @@ __device__ void search_core( if (threadIdx.x == 0) { *terminate_flag = 0; } } - topk_by_bitonic_sort_and_merge( - result_distances_buffer, - result_indices_buffer, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - search_width * graph_degree, - topk_ws, - (iter == 0), - multi_warps_1, - multi_warps_2); + topk_by_bitonic_sort_and_merge(result_distances_buffer, + result_indices_buffer, + max_itopk, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + search_width * graph_degree, + topk_ws, + (iter == 0), + multi_warps_1, + multi_warps_2); __syncthreads(); _CLK_REC(clk_topk); } else { _CLK_START(); // topk with radix block sort - topk_by_radix_sort{}( - internal_topk, - gridDim.x, - result_buffer_size, - reinterpret_cast(result_distances_buffer), - result_indices_buffer, - reinterpret_cast(result_distances_buffer), - result_indices_buffer, - nullptr, - topk_ws, - true, - smem_work_ptr); + topk_by_radix_sort{}(max_itopk, + (max_itopk <= 64) ? 32 : max_itopk / 4, + internal_topk, + gridDim.x, + result_buffer_size, + reinterpret_cast(result_distances_buffer), + result_indices_buffer, + reinterpret_cast(result_distances_buffer), + result_indices_buffer, + nullptr, + topk_ws, + true, + smem_work_ptr); _CLK_REC(clk_topk); // reset small-hash table @@ -877,18 +884,18 @@ __device__ void search_core( if (top_k > internal_topk || result_indices_buffer[top_k - 1] == invalid_index) { __syncthreads(); const unsigned multi_warps_1 = ((blockDim.x >= 64) && (MAX_CANDIDATES > 128)) ? 1 : 0; - const unsigned multi_warps_2 = ((blockDim.x >= 64) && (MAX_ITOPK > 256)) ? 1 : 0; - topk_by_bitonic_sort_and_merge( - result_distances_buffer, - result_indices_buffer, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - search_width * graph_degree, - topk_ws, - (iter == 0), - multi_warps_1, - multi_warps_2); + const unsigned multi_warps_2 = ((blockDim.x >= 64) && (max_itopk > 256)) ? 1 : 0; + topk_by_bitonic_sort_and_merge(result_distances_buffer, + result_indices_buffer, + max_itopk, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + search_width * graph_degree, + topk_ws, + (iter == 0), + multi_warps_1, + multi_warps_2); } __syncthreads(); } @@ -958,8 +965,7 @@ __device__ void search_core( #endif } -template bool return (h != kWaitForWork) && (h != kNoMoreWork); } -template auto dispatch_kernel = []() { if constexpr (Persistent) { - return search_kernel_p; } else { - return search_kernel struct search_kernel_config { using kernel_t = decltype(dispatch_kernel static auto choose_search_kernel(unsigned itopk_size) -> kernel_t { - if (itopk_size <= 64) { + if (itopk_size <= 512) { return dispatch_kernel; - } else if (itopk_size <= 128) { - return dispatch_kernel; - } else if (itopk_size <= 256) { - return dispatch_kernel; - } else if (itopk_size <= 512) { - return dispatch_kernel; } - THROW("No kernel for parametels itopk_size %u, max_candidates %u", itopk_size, MAX_CANDIDATES); + THROW("No kernel for parameters itopk_size %u, max_candidates %u", itopk_size, MAX_CANDIDATES); } static auto choose_itopk_and_mx_candidates(unsigned itopk_size, @@ -1285,25 +1263,14 @@ struct search_kernel_config { } else { // Radix-based topk is used constexpr unsigned max_candidates = 32; // to avoid build failure - if (itopk_size <= 256) { - return dispatch_kernel; - } else if (itopk_size <= 512) { - return dispatch_kernel; - } + return dispatch_kernel; } - THROW("No kernel for parametels itopk_size %u, num_itopk_candidates %u", + THROW("No kernel for parameters itopk_size %u, num_itopk_candidates %u", itopk_size, num_itopk_candidates); } @@ -1812,6 +1779,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b uint32_t num_random_samplings, uint64_t rand_xor_mask, uint32_t num_seeds, + uint32_t max_itopk, size_t itopk_size, size_t search_width, size_t min_iterations, @@ -1840,6 +1808,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b uint32_t num_random_samplings, uint64_t rand_xor_mask, uint32_t num_seeds, + uint32_t max_itopk, size_t itopk_size, size_t search_width, size_t min_iterations, @@ -1868,6 +1837,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b num_random_samplings, rand_xor_mask, num_seeds, + max_itopk, itopk_size, search_width, min_iterations, @@ -1938,6 +1908,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b &dev_seed_ptr, &num_seeds, &hashmap_ptr, // visited_hashmap_ptr: [num_queries, 1 << hash_bitlen] + &max_itopk, &itopk_size, &search_width, &min_iterations, @@ -2139,6 +2110,19 @@ void select_and_run( const SourceIndexT* source_indices_ptr = source_indices.has_value() ? source_indices->data_handle() : nullptr; + uint32_t max_itopk{}; + if (ps.itopk_size <= 64) { + max_itopk = 64; + } else if (ps.itopk_size <= 128) { + max_itopk = 128; + } else if (ps.itopk_size <= 256) { + max_itopk = 256; + } else if (ps.itopk_size <= 512) { + max_itopk = 512; + } else { + THROW("No kernel for parameter itopk_size %u", ps.itopk_size); + } + if (ps.persistent) { using runner_type = persistent_runner_t; @@ -2159,6 +2143,7 @@ control is returned in this thread (in persistent_runner_t constructor), so we'r ps.num_random_samplings, ps.rand_xor_mask, num_seeds, + max_itopk, ps.itopk_size, ps.search_width, ps.min_iterations, @@ -2190,6 +2175,7 @@ control is returned in this thread (in persistent_runner_t constructor), so we'r dev_seed_ptr, num_seeds, hashmap_ptr, + max_itopk, ps.itopk_size, ps.search_width, ps.min_iterations, diff --git a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh index 68aab54053..da571129a6 100644 --- a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh @@ -9,19 +9,18 @@ namespace cuvs::neighbors::cagra::detail { namespace single_cta_search { -template struct topk_by_radix_sort_base { - static constexpr std::uint32_t smem_size = MAX_INTERNAL_TOPK * 2 + 2048 + 8; - static constexpr std::uint32_t state_bit_lenght = 0; + static constexpr std::uint32_t state_bit_length = 0; static constexpr std::uint32_t vecLen = 2; // TODO + + static constexpr uint32_t smem_size(uint32_t max_itopk) { return max_itopk * 2 + 2048 + 8; } }; -template -struct topk_by_radix_sort : topk_by_radix_sort_base {}; -template -struct topk_by_radix_sort> - : topk_by_radix_sort_base { - RAFT_DEVICE_INLINE_FUNCTION void operator()(uint32_t topk, +template +struct topk_by_radix_sort : topk_by_radix_sort_base { + RAFT_DEVICE_INLINE_FUNCTION void operator()(uint32_t max_topk, + uint32_t num_sort_threads, + uint32_t topk, uint32_t batch_size, uint32_t len_x, const uint32_t* _x, @@ -33,48 +32,26 @@ struct topk_by_radix_sort= V / 4); std::uint8_t* const state = reinterpret_cast(work); - topk_cta_11_core::state_bit_lenght, - topk_by_radix_sort_base::vecLen, - 64, - 32, - IdxT>(topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); + topk_cta_11_core(max_topk, + num_sort_threads, + topk, + len_x, + _x, + _in_vals, + _y, + _out_vals, + state, + _hints, + sort, + _smem); } }; -#define TOP_FUNC_PARTIAL_SPECIALIZATION(V) \ - template \ - struct topk_by_radix_sort< \ - MAX_INTERNAL_TOPK, \ - IdxT, \ - std::enable_if_t<((MAX_INTERNAL_TOPK <= V) && (2 * MAX_INTERNAL_TOPK > V))>> \ - : topk_by_radix_sort_base { \ - RAFT_DEVICE_INLINE_FUNCTION void operator()(uint32_t topk, \ - uint32_t batch_size, \ - uint32_t len_x, \ - const uint32_t* _x, \ - const IdxT* _in_vals, \ - uint32_t* _y, \ - IdxT* _out_vals, \ - uint32_t* work, \ - uint32_t* _hints, \ - bool sort, \ - uint32_t* _smem) \ - { \ - assert(blockDim.x >= V / 4); \ - std::uint8_t* state = (std::uint8_t*)work; \ - topk_cta_11_core::state_bit_lenght, \ - topk_by_radix_sort_base::vecLen, \ - V, \ - V / 4, \ - IdxT>( \ - topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); \ - } \ - }; -TOP_FUNC_PARTIAL_SPECIALIZATION(128); -TOP_FUNC_PARTIAL_SPECIALIZATION(256); -TOP_FUNC_PARTIAL_SPECIALIZATION(512); -TOP_FUNC_PARTIAL_SPECIALIZATION(1024); - } // namespace single_cta_search } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk.cu b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk.cu index b8c343414e..41de68cede 100644 --- a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk.cu +++ b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk.cu @@ -60,6 +60,8 @@ void _cuann_find_topk(uint32_t topK, dim3 blocks(sizeBatch, 1, 1); void (*cta_kernel)(uint32_t, + uint32_t, + uint32_t, uint32_t, uint32_t, const uint32_t*, @@ -74,69 +76,73 @@ void _cuann_find_topk(uint32_t topK, uint32_t*, bool) = nullptr; - // V:vecLen, K:maxTopk, T:numSortThreads -#define SET_KERNEL_VKT(V, K, T, ValT) \ - do { \ - assert(numThreads >= T); \ - assert((K % T) == 0); \ - assert((K / T) <= 4); \ - cta_kernel = kern_topk_cta_11; \ - } while (0) - - // V: vecLen -#define SET_KERNEL_V(V, ValT) \ - do { \ - if (topK <= 32) { \ - SET_KERNEL_VKT(V, 32, 32, ValT); \ - } else if (topK <= 64) { \ - SET_KERNEL_VKT(V, 64, 32, ValT); \ - } else if (topK <= 96) { \ - SET_KERNEL_VKT(V, 96, 32, ValT); \ - } else if (topK <= 128) { \ - SET_KERNEL_VKT(V, 128, 32, ValT); \ - } else if (topK <= 192) { \ - SET_KERNEL_VKT(V, 192, 64, ValT); \ - } else if (topK <= 256) { \ - SET_KERNEL_VKT(V, 256, 64, ValT); \ - } else if (topK <= 384) { \ - SET_KERNEL_VKT(V, 384, 128, ValT); \ - } else if (topK <= 512) { \ - SET_KERNEL_VKT(V, 512, 128, ValT); \ - } else if (topK <= 768) { \ - SET_KERNEL_VKT(V, 768, 256, ValT); \ - } else if (topK <= 1024) { \ - SET_KERNEL_VKT(V, 1024, 256, ValT); \ - } \ - /* else if (topK <= 1536) { SET_KERNEL_VKT(V, 1536, 512); } */ \ - /* else if (topK <= 2048) { SET_KERNEL_VKT(V, 2048, 512); } */ \ - /* else if (topK <= 3072) { SET_KERNEL_VKT(V, 3072, 1024); } */ \ - /* else if (topK <= 4096) { SET_KERNEL_VKT(V, 4096, 1024); } */ \ - else { \ - RAFT_FAIL("topk must be lower than or equal to 1024"); \ - } \ - } while (0) - - int _vecLen = _get_vecLen(ldIK, 2); + int _vecLen = _get_vecLen(ldIK, 2); + constexpr int maxTopkPerThread = 4; + if (_vecLen == 2) { - SET_KERNEL_V(2, ValT); + cta_kernel = kern_topk_cta_11; } else if (_vecLen == 1) { - SET_KERNEL_V(1, ValT); + cta_kernel = kern_topk_cta_11; + } + + int max_topk{}; + int num_sort_threads{}; + if (topK <= 32) { + max_topk = 32; + num_sort_threads = 32; + } else if (topK <= 64) { + max_topk = 64; + num_sort_threads = 32; + } else if (topK <= 96) { + max_topk = 96; + num_sort_threads = 32; + } else if (topK <= 128) { + max_topk = 128; + num_sort_threads = 32; + } else if (topK <= 192) { + max_topk = 192; + num_sort_threads = 64; + } else if (topK <= 256) { + max_topk = 256; + num_sort_threads = 64; + } else if (topK <= 384) { + max_topk = 384; + num_sort_threads = 128; + } else if (topK <= 512) { + max_topk = 512; + num_sort_threads = 128; + } else if (topK <= 768) { + max_topk = 768; + num_sort_threads = 256; + } else if (topK <= 1024) { + max_topk = 1024; + num_sort_threads = 256; + } else { + RAFT_FAIL("topK must be lower than or equal to 1024"); } - cta_kernel<<>>(topK, - sizeBatch, - numElements, - (const uint32_t*)inputKeys, - ldIK, - inputVals, - ldIV, - (uint32_t*)outputKeys, - ldOK, - outputVals, - ldOV, - state, - hints, - sort); + assert(max_topk % num_sort_threads == 0); + assert(max_topk / num_sort_threads <= maxTopkPerThread); + + const size_t smem_len = 2 * max_topk + 2048 + 8; + assert(max_topk * (1 + utils::size_of() / utils::size_of()) <= smem_len); + const size_t smem_size = smem_len * sizeof(uint32_t); + cta_kernel<<>>(max_topk, + num_sort_threads, + topK, + sizeBatch, + numElements, + (const uint32_t*)inputKeys, + ldIK, + inputVals, + ldIV, + (uint32_t*)outputKeys, + ldOK, + outputVals, + ldOV, + state, + hints, + sort); return; } diff --git a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh index 7eb6a8a1a0..6e07587aa3 100644 --- a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh @@ -599,9 +599,11 @@ RAFT_INLINE_FUNCTION constexpr uint32_t get_state_size(uint32_t len_x) return 0; } -// -template -RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, +// max_topk / num_sort_threads should never exceed maxTopkPerThread +template +RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, + uint32_t num_sort_threads, + uint32_t topk, uint32_t len_x, const uint32_t* _x, // [size_batch, ld_x,] const ValT* _in_vals, // [size_batch, ld_iv,] @@ -613,9 +615,9 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, uint32_t* _smem) { uint32_t* const smem_out_vals = _smem; - uint32_t* const hist = &(_smem[2 * maxTopk]); - uint32_t* const best_index = &(_smem[2 * maxTopk + 2048]); - uint32_t* const best_csum = &(_smem[2 * maxTopk + 2048 + 3]); + uint32_t* const hist = &(_smem[2 * max_topk]); + uint32_t* const best_index = &(_smem[2 * max_topk + 2048]); + uint32_t* const best_csum = &(_smem[2 * max_topk + 2048 + 3]); const uint32_t num_threads = blockDim.x; const uint32_t thread_id = threadIdx.x; @@ -631,11 +633,11 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, const uint32_t hint = (_hint == NULL ? ~0u : *_hint); // Initialize shared memory - for (int i = 2 * maxTopk + thread_id; i < 2 * maxTopk + 2048 + 8; i += num_threads) { + for (int i = 2 * max_topk + thread_id; i < 2 * max_topk + 2048 + 8; i += num_threads) { _smem[i] = 0; } - uint32_t* const output_count = &(_smem[2 * maxTopk + 2048 + 6]); - uint32_t* const output_count_eq = &(_smem[2 * maxTopk + 2048 + 7]); + uint32_t* const output_count = &(_smem[2 * max_topk + 2048 + 6]); + uint32_t* const output_count_eq = &(_smem[2 * max_topk + 2048 + 7]); uint32_t threshold = 0; uint32_t nx_below_threshold = 0; __syncthreads(); @@ -721,14 +723,18 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, return; } - constexpr int numTopkPerThread = maxTopk / numSortThreads; - float my_keys[numTopkPerThread]; - ValT my_vals[numTopkPerThread]; + const int num_topk_per_thread = max_topk / num_sort_threads; + static_assert( + maxTopkPerThread <= + 4); // we will end-up wasting too much stack memory if maxTopkPerThread >> num_topk_per_thread + assert(num_topk_per_thread <= maxTokenPerThread); + float my_keys[maxTopkPerThread]; + ValT my_vals[maxTopkPerThread]; // Read keys and values to registers - if (thread_id < numSortThreads) { - for (int i = 0; i < numTopkPerThread; i++) { - const int k = thread_id + (numSortThreads * i); + if (thread_id < num_sort_threads) { + for (int i = 0; i < num_topk_per_thread; i++) { + const int k = thread_id + (num_sort_threads * i); if (k < topk) { const int j = smem_out_vals[k]; my_keys[i] = ((float*)x)[j]; @@ -747,21 +753,21 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, uint32_t mask = 1; // Sorting by thread - if (thread_id < numSortThreads) { + if (thread_id < num_sort_threads) { const bool ascending = ((thread_id & mask) == 0); - if (numTopkPerThread == 3) { + if (num_topk_per_thread == 3) { swap_if_needed(my_keys[0], my_keys[1], my_vals[0], my_vals[1], ascending); swap_if_needed(my_keys[0], my_keys[2], my_vals[0], my_vals[2], ascending); swap_if_needed(my_keys[1], my_keys[2], my_vals[1], my_vals[2], ascending); } else { - for (int j = 0; j < numTopkPerThread / 2; j += 1) { + for (int j = 0; j < num_topk_per_thread / 2; j += 1) { #pragma unroll - for (int i = 0; i < numTopkPerThread; i += 2) { + for (int i = 0; i < num_topk_per_thread; i += 2) { swap_if_needed( my_keys[i], my_keys[i + 1], my_vals[i], my_vals[i + 1], ascending); } #pragma unroll - for (int i = 1; i < numTopkPerThread - 1; i += 2) { + for (int i = 1; i < num_topk_per_thread - 1; i += 2) { swap_if_needed( my_keys[i], my_keys[i + 1], my_vals[i], my_vals[i + 1], ascending); } @@ -770,38 +776,38 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, } // Bitonic Sorting - while (mask < numSortThreads) { + while (mask < num_sort_threads) { uint32_t next_mask = mask << 1; for (uint32_t curr_mask = mask; curr_mask > 0; curr_mask >>= 1) { const bool ascending = ((thread_id & curr_mask) == 0) == ((thread_id & next_mask) == 0); if (curr_mask >= 32) { // inter warp - ValT* const smem_vals = reinterpret_cast(_smem); // [maxTopk] - float* const smem_keys = - reinterpret_cast(smem_vals + maxTopk); // [numTopkPerThread, numSortThreads] + ValT* const smem_vals = reinterpret_cast(_smem); // [max_topk] + float* const smem_keys = reinterpret_cast( + smem_vals + max_topk); // [num_topk_per_thread, num_sort_threads] __syncthreads(); - if (thread_id < numSortThreads) { + if (thread_id < num_sort_threads) { #pragma unroll - for (int i = 0; i < numTopkPerThread; i++) { - smem_keys[thread_id + (numSortThreads * i)] = my_keys[i]; - smem_vals[thread_id + (numSortThreads * i)] = my_vals[i]; + for (int i = 0; i < num_topk_per_thread; i++) { + smem_keys[thread_id + (num_sort_threads * i)] = my_keys[i]; + smem_vals[thread_id + (num_sort_threads * i)] = my_vals[i]; } } __syncthreads(); - if (thread_id < numSortThreads) { + if (thread_id < num_sort_threads) { #pragma unroll - for (int i = 0; i < numTopkPerThread; i++) { - float opp_key = smem_keys[(thread_id ^ curr_mask) + (numSortThreads * i)]; - ValT opp_val = smem_vals[(thread_id ^ curr_mask) + (numSortThreads * i)]; + for (int i = 0; i < num_topk_per_thread; i++) { + float opp_key = smem_keys[(thread_id ^ curr_mask) + (num_sort_threads * i)]; + ValT opp_val = smem_vals[(thread_id ^ curr_mask) + (num_sort_threads * i)]; swap_if_needed(my_keys[i], opp_key, my_vals[i], opp_val, ascending); } } } else { // intra warp - if (thread_id < numSortThreads) { + if (thread_id < num_sort_threads) { #pragma unroll - for (int i = 0; i < numTopkPerThread; i++) { + for (int i = 0; i < num_topk_per_thread; i++) { float opp_key = __shfl_xor_sync(0xffffffff, my_keys[i], curr_mask); ValT opp_val = __shfl_xor_sync(0xffffffff, my_vals[i], curr_mask); swap_if_needed(my_keys[i], opp_key, my_vals[i], opp_val, ascending); @@ -810,17 +816,17 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, } } - if (thread_id < numSortThreads) { + if (thread_id < num_sort_threads) { const bool ascending = ((thread_id & next_mask) == 0); - if (numTopkPerThread == 3) { + if (num_topk_per_thread == 3) { swap_if_needed(my_keys[0], my_keys[1], my_vals[0], my_vals[1], ascending); swap_if_needed(my_keys[0], my_keys[2], my_vals[0], my_vals[2], ascending); swap_if_needed(my_keys[1], my_keys[2], my_vals[1], my_vals[2], ascending); } else { #pragma unroll - for (uint32_t curr_mask = numTopkPerThread / 2; curr_mask > 0; curr_mask >>= 1) { + for (uint32_t curr_mask = num_topk_per_thread / 2; curr_mask > 0; curr_mask >>= 1) { #pragma unroll - for (int i = 0; i < numTopkPerThread; i++) { + for (int i = 0; i < num_topk_per_thread; i++) { const int j = i ^ curr_mask; if (i > j) continue; swap_if_needed(my_keys[i], my_keys[j], my_vals[i], my_vals[j], ascending); @@ -832,9 +838,9 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, } // Write sorted keys and values - if (thread_id < numSortThreads) { - for (int i = 0; i < numTopkPerThread; i++) { - const int k = i + (numTopkPerThread * thread_id); + if (thread_id < num_sort_threads) { + for (int i = 0; i < num_topk_per_thread; i++) { + const int k = i + (num_topk_per_thread * thread_id); if (k < topk) { if (y) { y[k] = reinterpret_cast(my_keys)[i]; } if (out_vals) { out_vals[k] = my_vals[i]; } @@ -862,9 +868,11 @@ int _get_vecLen(uint32_t maxSamples, int maxVecLen = MAX_VEC_LENGTH) } } // unnamed namespace -template +template __launch_bounds__(1024, 1) RAFT_KERNEL - kern_topk_cta_11(uint32_t topk, + kern_topk_cta_11(uint32_t max_topk, + uint32_t num_sort_threads, + uint32_t topk, uint32_t size_batch, uint32_t len_x, const uint32_t* _x, // [size_batch, ld_x,] @@ -882,12 +890,11 @@ __launch_bounds__(1024, 1) RAFT_KERNEL const uint32_t i_batch = blockIdx.x; if (i_batch >= size_batch) return; - constexpr uint32_t smem_len = 2 * maxTopk + 2048 + 8; - static_assert(maxTopk * (1 + utils::size_of() / utils::size_of()) <= smem_len, - "maxTopk * sizeof(ValT) must be smaller or equal to 8192 byte"); - __shared__ uint32_t _smem[smem_len]; + extern __shared__ uint32_t _smem[]; - topk_cta_11_core( + topk_cta_11_core( + max_topk, + num_sort_threads, topk, len_x, (_x == NULL ? NULL : _x + i_batch * ld_x), From 444b946d530ba326398bc4cab6e2a14c8ca45964 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 5 Nov 2025 01:00:04 -0800 Subject: [PATCH 02/25] convert max_candidates to a runtime parameter --- cpp/src/neighbors/detail/cagra/bitonic.hpp | 43 +-- .../cagra/search_single_cta_kernel-inl.cuh | 292 +++++++++--------- 2 files changed, 154 insertions(+), 181 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/bitonic.hpp b/cpp/src/neighbors/detail/cagra/bitonic.hpp index 6a003face9..85a928d3bd 100644 --- a/cpp/src/neighbors/detail/cagra/bitonic.hpp +++ b/cpp/src/neighbors/detail/cagra/bitonic.hpp @@ -18,8 +18,12 @@ template RAFT_DEVICE_INLINE_FUNCTION void swap_if_needed(K& k0, V& v0, K& k1, V& v1, const bool asc) { if ((k0 != k1) && ((k0 < k1) != asc)) { - cuda::std::swap(k0, k1); - cuda::std::swap(v0, v1); + const auto tmp_k = k0; + k0 = k1; + k1 = tmp_k; + const auto tmp_v = v0; + v0 = v1; + v1 = tmp_v; } } @@ -79,38 +83,6 @@ struct warp_merge_core_n { } }; -template -struct warp_merge_core_3 { - RAFT_DEVICE_INLINE_FUNCTION void operator()(K* ks, - V* vs, - const std::uint32_t range, - const bool asc) - { - constexpr unsigned N = 3; - const auto lane_id = threadIdx.x % warp_size; - - if (range == 1) { - const auto p = ((lane_id & 1) == 0); - swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); - swap_if_needed(ks[1], vs[1], ks[2], vs[2], p); - swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); - } else { - const std::uint32_t b = range; - for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { - const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); -#pragma unroll - for (std::uint32_t i = 0; i < N; i++) { - swap_if_needed(ks[i], vs[i], c, p); - } - } - const auto p = ((lane_id & b) == 0); - swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); - swap_if_needed(ks[1], vs[1], ks[2], vs[2], p); - swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); - } - } -}; - template struct warp_merge_core_2 { RAFT_DEVICE_INLINE_FUNCTION void operator()(K* ks, @@ -165,8 +137,6 @@ RAFT_DEVICE_INLINE_FUNCTION void warp_merge( detail::warp_merge_core_1{}(ks, vs, range, asc); } else if (n == 2) { detail::warp_merge_core_2{}(ks, vs, range, asc); - } else if (n == 3) { - detail::warp_merge_core_3{}(ks, vs, range, asc); } else { detail::warp_merge_core_n{}(ks, vs, n, range, asc); } @@ -175,6 +145,7 @@ RAFT_DEVICE_INLINE_FUNCTION void warp_merge( template RAFT_DEVICE_INLINE_FUNCTION void warp_sort(K* ks, V* vs, unsigned n, const bool asc = true) { +#pragma unroll 1 for (std::uint32_t range = 1; range <= warp_size; range <<= 1) { warp_merge(ks, vs, n, range, asc); } diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 7e2cbc55bb..cd80c68477 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -99,21 +99,27 @@ RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parents(std::uint32_t* const termin if (threadIdx.x == 0 && (num_new_parents == 0)) { *terminate_flag = 1; } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] + const std::uint32_t max_candidates, const std::uint32_t num_candidates, const std::uint32_t num_itopk, - unsigned MULTI_WARPS = 0) + bool multi_warps = false) { const unsigned lane_id = threadIdx.x % raft::warp_size(); const unsigned warp_id = threadIdx.x / raft::warp_size(); - if (MULTI_WARPS == 0) { + assert(max_candidates <= 256); + constexpr unsigned MAX_N = + 8; // if MAX_N >> N, we may have negative performance impact, if this is significant, we may + // get memory space from dynamically sized shared memory. + float key[MAX_N]; + IdxT val[MAX_N]; + if (multi_warps == false) { if (warp_id > 0) { return; } - constexpr unsigned N = (MAX_CANDIDATES + (raft::warp_size() - 1)) / raft::warp_size(); - float key[N]; - IdxT val[N]; + const unsigned N = (max_candidates + (raft::warp_size() - 1)) / raft::warp_size(); + assert(N <= MAX_N); /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { unsigned j = lane_id + (raft::warp_size() * i); @@ -137,10 +143,9 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } else { // Use two warps (64 threads) - constexpr unsigned max_candidates_per_warp = (MAX_CANDIDATES + 1) / 2; - constexpr unsigned N = (max_candidates_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); - float key[N]; - IdxT val[N]; + const unsigned max_candidates_per_warp = (max_candidates + 1) / 2; + const unsigned N = (max_candidates_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); + assert(N <= MAX_N); if (warp_id < 2) { /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { @@ -175,7 +180,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( unsigned jl = (N * lane_id) + i; unsigned kl = max_candidates_per_warp - 1 - jl; unsigned j = jl + (max_candidates_per_warp * warp_id); - unsigned k = MAX_CANDIDATES - 1 - j; + unsigned k = max_candidates - 1 - j; if (j >= num_candidates || k >= num_candidates || kl >= num_itopk) continue; float temp_key = candidate_distances[device::swizzling(k)]; if (key[i] == temp_key) continue; @@ -214,18 +219,18 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( const std::uint32_t num_candidates, std::uint32_t* work_buf, const bool first, - unsigned MULTI_WARPS = 0) + bool multi_warps = false) { const unsigned lane_id = threadIdx.x % raft::warp_size(); const unsigned warp_id = threadIdx.x / raft::warp_size(); + assert(max_itopk <= 1024); constexpr unsigned MAX_N = - 32; // if MAX_N >> N, we may have occupancy reduction as this inflates per-thread stack size. - // If this has significant performance impact, we may get memory space for key & val from - // dynamically sized shared memory. + 32; // if MAX_N >> N, we may have negative performance impact, if this is significant, we may + // get memory space from dynamically sized shared memory. float key[MAX_N]; IdxT val[MAX_N]; - if (MULTI_WARPS == 0) { + if (multi_warps == false) { if (warp_id > 0) { return; } const unsigned N = (max_itopk + (raft::warp_size() - 1)) / raft::warp_size(); assert(N <= MAX_N); @@ -416,8 +421,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] @@ -425,15 +429,20 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( const std::uint32_t num_itopk, float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] + const std::uint32_t max_candidates, const std::uint32_t num_candidates, std::uint32_t* work_buf, const bool first, - const unsigned MULTI_WARPS_1, - const unsigned MULTI_WARPS_2) + const bool multi_warps_1, + const bool multi_warps_2) { // The results in candidate_distances/indices are sorted by bitonic sort. - topk_by_bitonic_sort_and_full( - candidate_distances, candidate_indices, num_candidates, num_itopk, MULTI_WARPS_1); + topk_by_bitonic_sort_and_full(candidate_distances, + candidate_indices, + max_candidates, + num_candidates, + num_itopk, + multi_warps_1); // The results sorted above are merged with the internal intermediate top-k // results so far using bitonic merge. @@ -446,7 +455,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( num_candidates, work_buf, first, - MULTI_WARPS_2); + multi_warps_2); } // This function move the invalid index element to the end of the itopk list. @@ -509,7 +518,6 @@ RAFT_DEVICE_INLINE_FUNCTION void hashmap_restore(INDEX_T* const hashmap_ptr, /** * @brief Search operation for a single query using a single thread block. * * - * @tparam MAX_CANDIDATES * @tparam TOPK_BY_BITONIC_SORT * @tparam DATASET_DESCRIPTOR_T * @tparam SAMPLE_FILTER_T @@ -539,8 +547,7 @@ RAFT_DEVICE_INLINE_FUNCTION void hashmap_restore(INDEX_T* const hashmap_ptr, * @param small_hash_reset_interval Interval for resetting the small hash. * @param query_id sequential id of the query in the batch */ -template @@ -559,6 +566,7 @@ __device__ void search_core( const uint32_t num_seeds, typename DATASET_DESCRIPTOR_T::INDEX_T* const visited_hashmap_ptr, // [num_queries, 1 << hash_bitlen] + const std::uint32_t max_candidates, const std::uint32_t max_itopk, const std::uint32_t internal_topk, const std::uint32_t search_width, @@ -671,10 +679,10 @@ __device__ void search_core( // batch size is small (short-latency), but it might not be always good // when batch size is large (high-throughput). // topk_by_bitonic_sort_and_merge() consists of two operations: - // if MAX_CANDIDATES is greater than 128, the first operation uses two warps; + // if max_candidates is greater than 128, the first operation uses two warps; // if max_itopk is greater than 256, the second operation used two warps. - const unsigned multi_warps_1 = ((blockDim.x >= 64) && (MAX_CANDIDATES > 128)) ? 1 : 0; - const unsigned multi_warps_2 = ((blockDim.x >= 64) && (max_itopk > 256)) ? 1 : 0; + const bool multi_warps_1 = ((blockDim.x >= 64) && (max_candidates > 128)) ? true : false; + const bool multi_warps_2 = ((blockDim.x >= 64) && (max_itopk > 256)) ? true : false; // reset small-hash table. if ((iter + 1) % small_hash_reset_interval == 0) { @@ -715,17 +723,18 @@ __device__ void search_core( if (threadIdx.x == 0) { *terminate_flag = 0; } } - topk_by_bitonic_sort_and_merge(result_distances_buffer, - result_indices_buffer, - max_itopk, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - search_width * graph_degree, - topk_ws, - (iter == 0), - multi_warps_1, - multi_warps_2); + topk_by_bitonic_sort_and_merge(result_distances_buffer, + result_indices_buffer, + max_itopk, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + max_candidates, + search_width * graph_degree, + topk_ws, + (iter == 0), + multi_warps_1, + multi_warps_2); __syncthreads(); _CLK_REC(clk_topk); } else { @@ -883,19 +892,20 @@ __device__ void search_core( // candidate list. if (top_k > internal_topk || result_indices_buffer[top_k - 1] == invalid_index) { __syncthreads(); - const unsigned multi_warps_1 = ((blockDim.x >= 64) && (MAX_CANDIDATES > 128)) ? 1 : 0; - const unsigned multi_warps_2 = ((blockDim.x >= 64) && (max_itopk > 256)) ? 1 : 0; - topk_by_bitonic_sort_and_merge(result_distances_buffer, - result_indices_buffer, - max_itopk, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - search_width * graph_degree, - topk_ws, - (iter == 0), - multi_warps_1, - multi_warps_2); + const bool multi_warps_1 = ((blockDim.x >= 64) && (max_candidates > 128)) ? true : false; + const bool multi_warps_2 = ((blockDim.x >= 64) && (max_itopk > 256)) ? true : false; + topk_by_bitonic_sort_and_merge(result_distances_buffer, + result_indices_buffer, + max_itopk, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + max_candidates, + search_width * graph_degree, + topk_ws, + (iter == 0), + multi_warps_1, + multi_warps_2); } __syncthreads(); } @@ -925,7 +935,7 @@ __device__ void search_core( for (std::uint32_t i = threadIdx.x; i < top_k; i += blockDim.x) { unsigned j = i + (top_k * query_id); unsigned ii = i; - if (TOPK_BY_BITONIC_SORT) { ii = device::swizzling(i); } + if constexpr (TOPK_BY_BITONIC_SORT) { ii = device::swizzling(i); } if (result_distances_ptr != nullptr) { result_distances_ptr[j] = result_distances_buffer[ii]; } constexpr INDEX_T index_msb_1_mask = utils::gen_index_msb_1_mask::value; @@ -965,8 +975,7 @@ __device__ void search_core( #endif } -template @@ -985,6 +994,7 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( const uint32_t num_seeds, typename DATASET_DESCRIPTOR_T::INDEX_T* const visited_hashmap_ptr, // [num_queries, 1 << hash_bitlen] + const std::uint32_t max_candidates, const std::uint32_t max_itopk, const std::uint32_t internal_topk, const std::uint32_t search_width, @@ -997,34 +1007,32 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( SAMPLE_FILTER_T sample_filter) { const auto query_id = blockIdx.y; - search_core(result_indices_ptr, - result_distances_ptr, - top_k, - dataset_desc, - queries_ptr, - knn_graph, - graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hashmap_ptr, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - sample_filter); + search_core( + result_indices_ptr, + result_distances_ptr, + top_k, + dataset_desc, + queries_ptr, + knn_graph, + graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + sample_filter); } // To make sure we avoid false sharing on both CPU and GPU, we enforce cache line size to the @@ -1086,8 +1094,7 @@ constexpr auto is_worker_busy(worker_handle_t::handle_t h) -> bool return (h != kWaitForWork) && (h != kNoMoreWork); } -template @@ -1105,6 +1112,7 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel_p( const uint32_t num_seeds, typename DATASET_DESCRIPTOR_T::INDEX_T* const visited_hashmap_ptr, // [num_queries, 1 << hash_bitlen] + const std::uint32_t max_candidates, const std::uint32_t max_itopk, const std::uint32_t internal_topk, const std::uint32_t search_width, @@ -1157,34 +1165,32 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel_p( auto query_id = worker_data.value.query_id; // work phase - search_core(result_indices_ptr, - result_distances_ptr, - top_k, - dataset_desc, - queries_ptr, - knn_graph, - graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hashmap_ptr, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - sample_filter); + search_core( + result_indices_ptr, + result_distances_ptr, + top_k, + dataset_desc, + queries_ptr, + knn_graph, + graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + sample_filter); // make sure all writes are visible even for the host // (e.g. when result buffers are in pinned memory) @@ -1202,24 +1208,18 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel_p( } template auto dispatch_kernel = []() { if constexpr (Persistent) { - return search_kernel_p; } else { - return search_kernel; + return search_kernel; } }(); @@ -1228,47 +1228,31 @@ template struct search_kernel_config { - using kernel_t = decltype(dispatch_kernel); + using kernel_t = + decltype(dispatch_kernel); - template + template static auto choose_search_kernel(unsigned itopk_size) -> kernel_t { if (itopk_size <= 512) { return dispatch_kernel; } - THROW("No kernel for parameters itopk_size %u, max_candidates %u", itopk_size, MAX_CANDIDATES); + THROW("No kernel for parameters itopk_size %u", itopk_size); } static auto choose_itopk_and_mx_candidates(unsigned itopk_size, unsigned num_itopk_candidates, unsigned block_size) -> kernel_t { - if (num_itopk_candidates <= 64) { - // use bitonic sort based topk - return choose_search_kernel<64, 1>(itopk_size); - } else if (num_itopk_candidates <= 128) { - return choose_search_kernel<128, 1>(itopk_size); - } else if (num_itopk_candidates <= 256) { - return choose_search_kernel<256, 1>(itopk_size); + if (num_itopk_candidates <= 256) { + return choose_search_kernel<1>(itopk_size); } else { // Radix-based topk is used - constexpr unsigned max_candidates = 32; // to avoid build failure - return dispatch_kernel; + return dispatch_kernel; } THROW("No kernel for parameters itopk_size %u, num_itopk_candidates %u", itopk_size, @@ -1770,6 +1754,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b std::reference_wrapper> dataset_desc, raft::device_matrix_view graph, const SourceIndexT* source_indices_ptr, + uint32_t max_candidates, uint32_t num_itopk_candidates, uint32_t block_size, // uint32_t smem_size, @@ -1799,6 +1784,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b std::reference_wrapper> dataset_desc, raft::device_matrix_view graph, const SourceIndexT* source_indices_ptr, + uint32_t max_candidates, uint32_t num_itopk_candidates, uint32_t block_size, // uint32_t smem_size, @@ -1828,6 +1814,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b param_hash(calculate_parameter_hash(dd_host, graph, source_indices_ptr, + max_candidates, num_itopk_candidates, block_size, smem_size, @@ -1908,6 +1895,7 @@ struct alignas(kCacheLineBytes) persistent_runner_t : public persistent_runner_b &dev_seed_ptr, &num_seeds, &hashmap_ptr, // visited_hashmap_ptr: [num_queries, 1 << hash_bitlen] + &max_candidates, &max_itopk, &itopk_size, &search_width, @@ -2123,6 +2111,18 @@ void select_and_run( THROW("No kernel for parameter itopk_size %u", ps.itopk_size); } + uint32_t max_candidates{}; + if (num_itopk_candidates <= 64) { + max_candidates = 64; + } else if (num_itopk_candidates <= 128) { + max_candidates = 128; + } else if (num_itopk_candidates <= 256) { + max_candidates = 256; + } else { + max_candidates = + 32; // irrelevant, radix based topk is used (see choose_itopk_and_max_candidates) + } + if (ps.persistent) { using runner_type = persistent_runner_t; @@ -2143,6 +2143,7 @@ control is returned in this thread (in persistent_runner_t constructor), so we'r ps.num_random_samplings, ps.rand_xor_mask, num_seeds, + max_candidates, max_itopk, ps.itopk_size, ps.search_width, @@ -2175,6 +2176,7 @@ control is returned in this thread (in persistent_runner_t constructor), so we'r dev_seed_ptr, num_seeds, hashmap_ptr, + max_candidates, max_itopk, ps.itopk_size, ps.search_width, From 54a1805bf44e44d20f81a1120fb4a32a079d3665 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 5 Nov 2025 01:25:35 -0800 Subject: [PATCH 03/25] convert max_elements to a runtime parameter --- .../cagra/search_multi_cta_kernel-inl.cuh | 55 +++++++++++-------- 1 file changed, 33 insertions(+), 22 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 7a50edc884..879c41487a 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -109,19 +109,25 @@ RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parent( } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort(float* distances, // [num_elements] INDEX_T* indices, // [num_elements] + const uint32_t max_elements, const uint32_t num_elements) { - const unsigned warp_id = threadIdx.x / 32; + const unsigned warp_id = threadIdx.x / raft::warp_size(); if (warp_id > 0) { return; } - const unsigned lane_id = threadIdx.x % 32; - constexpr unsigned N = (MAX_ELEMENTS + 31) / 32; - float key[N]; - INDEX_T val[N]; + const unsigned lane_id = threadIdx.x % raft::warp_size(); + assert(max_elements <= 256); + constexpr unsigned MAX_N = + 8; // if MAX_N >> N, we may have negative performance impact, if this is significant, we may + // get memory space from dynamically sized shared memory + const unsigned N = (max_elements + (raft::warp_size() - 1)) / raft::warp_size(); + assert(N <= MAX_N); + float key[MAX_N]; + INDEX_T val[MAX_N]; for (unsigned i = 0; i < N; i++) { - unsigned j = lane_id + (32 * i); + unsigned j = lane_id + (raft::warp_size() * i); if (j < num_elements) { key[i] = distances[j]; val[i] = indices[j]; @@ -145,10 +151,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort(float* distances, // [num // // multiple CTAs per single query // -template +template RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( typename DATASET_DESCRIPTOR_T::INDEX_T* const result_indices_ptr, // [num_queries, num_cta_per_query, itopk_size] @@ -157,6 +160,7 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( const DATASET_DESCRIPTOR_T* dataset_desc, const typename DATASET_DESCRIPTOR_T::DATA_T* const queries_ptr, // [num_queries, dataset_dim] const typename DATASET_DESCRIPTOR_T::INDEX_T* const knn_graph, // [dataset_size, graph_degree] + const uint32_t max_elements, const uint32_t graph_degree, const SourceIndexT* source_indices_ptr, // [num_queries, search_width] const unsigned num_distilation, @@ -211,7 +215,7 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( // |<--- result_buffer_size_32 --->| const auto result_buffer_size = itopk_size + graph_degree; const auto result_buffer_size_32 = raft::round_up_safe(result_buffer_size, 32); - assert(result_buffer_size_32 <= MAX_ELEMENTS); + assert(result_buffer_size_32 <= max_elements); // Set smem working buffer for the distance calculation dataset_desc = dataset_desc->setup_workspace(smem, queries_ptr, query_id); @@ -268,8 +272,8 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( _CLK_START(); if (threadIdx.x < 32) { // [1st warp] Topk with bitonic sort - topk_by_bitonic_sort( - result_distances_buffer, result_indices_buffer, result_buffer_size_32); + topk_by_bitonic_sort( + result_distances_buffer, result_indices_buffer, max_elements, result_buffer_size_32); } __syncthreads(); _CLK_REC(clk_topk); @@ -487,17 +491,12 @@ struct search_kernel_config { // Search kernel function type. Note that the actual values for the template value // parameters do not matter, because they are not part of the function signature. The // second to fourth value parameters will be selected by the choose_* functions below. - using kernel_t = - decltype(&search_kernel<128, DATASET_DESCRIPTOR_T, SourceIndexT, SAMPLE_FILTER_T>); + using kernel_t = decltype(&search_kernel); static auto choose_buffer_size(unsigned result_buffer_size, unsigned block_size) -> kernel_t { - if (result_buffer_size <= 64) { - return search_kernel<64, DATASET_DESCRIPTOR_T, SourceIndexT, SAMPLE_FILTER_T>; - } else if (result_buffer_size <= 128) { - return search_kernel<128, DATASET_DESCRIPTOR_T, SourceIndexT, SAMPLE_FILTER_T>; - } else if (result_buffer_size <= 256) { - return search_kernel<256, DATASET_DESCRIPTOR_T, SourceIndexT, SAMPLE_FILTER_T>; + if (result_buffer_size <= 256) { + return search_kernel; } THROW("Result buffer size %u larger than max buffer size %u", result_buffer_size, 256); } @@ -536,6 +535,17 @@ void select_and_run(const dataset_descriptor_host& dat SourceIndexT, SampleFilterT>::choose_buffer_size(result_buffer_size, block_size); + uint32_t max_elements{}; + if (result_buffer_size <= 64) { + max_elements = 64; + } else if (result_buffer_size <= 128) { + max_elements = 128; + } else if (result_buffer_size <= 256) { + max_elements = 256; + } else { + THROW("Result buffer size %u larger than max buffer size %u", result_buffer_size, 256); + } + RAFT_CUDA_TRY( cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); // Initialize hash table @@ -560,6 +570,7 @@ void select_and_run(const dataset_descriptor_host& dat dataset_desc.dev_ptr(stream), queries_ptr, graph.data_handle(), + max_elements, graph.extent(1), source_indices_ptr, ps.num_random_samplings, From aa5c116a40bf1758b284504a44b356822b27bb36 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Mon, 10 Nov 2025 20:19:39 -0800 Subject: [PATCH 04/25] tighter bound on array size --- .../neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index cd80c68477..ed6f18daa6 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -224,9 +224,9 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( const unsigned lane_id = threadIdx.x % raft::warp_size(); const unsigned warp_id = threadIdx.x / raft::warp_size(); - assert(max_itopk <= 1024); + assert(max_itopk <= 512); constexpr unsigned MAX_N = - 32; // if MAX_N >> N, we may have negative performance impact, if this is significant, we may + 16; // if MAX_N >> N, we may have negative performance impact, if this is significant, we may // get memory space from dynamically sized shared memory. float key[MAX_N]; IdxT val[MAX_N]; From 178459a6d83cc97543e733c32b0d3b3a92137f3b Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 12 Nov 2025 09:54:28 -0800 Subject: [PATCH 05/25] undo most of the changes in bitonic.hpp except for the less unrolling on successive warp_merge calls when N is large, we now handle run-time branch in a higher level --- cpp/src/neighbors/detail/cagra/bitonic.hpp | 194 +++++++++++++++------ 1 file changed, 141 insertions(+), 53 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/bitonic.hpp b/cpp/src/neighbors/detail/cagra/bitonic.hpp index 85a928d3bd..58cd83a5b3 100644 --- a/cpp/src/neighbors/detail/cagra/bitonic.hpp +++ b/cpp/src/neighbors/detail/cagra/bitonic.hpp @@ -41,52 +41,139 @@ RAFT_DEVICE_INLINE_FUNCTION void swap_if_needed(K& k0, } } -template -struct warp_merge_core_n { - RAFT_DEVICE_INLINE_FUNCTION void operator()( - K* ks, V* vs, unsigned n, const std::uint32_t range, const bool asc) +template +struct warp_merge_core { + RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[N], + V v[N], + const std::uint32_t range, + const bool asc) { const auto lane_id = threadIdx.x % warp_size; if (range == 1) { - for (std::uint32_t b = 2; b <= n; b <<= 1) { + for (std::uint32_t b = 2; b <= N; b <<= 1) { for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { #pragma unroll - for (std::uint32_t i = 0; i < n; i++) { + for (std::uint32_t i = 0; i < N; i++) { std::uint32_t j = i ^ c; if (i >= j) continue; - const auto line_id = i + (n * lane_id); + const auto line_id = i + (N * lane_id); const auto p = static_cast(line_id & b) == static_cast(line_id & c); - swap_if_needed(ks[i], vs[i], ks[j], vs[j], p); + swap_if_needed(k[i], v[i], k[j], v[j], p); } } } - } else { - const std::uint32_t b = range; - for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { - const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); + return; + } + + const std::uint32_t b = range; + for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { + const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); #pragma unroll - for (std::uint32_t i = 0; i < n; i++) { - swap_if_needed(ks[i], vs[i], c, p); - } + for (std::uint32_t i = 0; i < N; i++) { + swap_if_needed(k[i], v[i], c, p); } - const auto p = ((lane_id & b) == 0); - for (std::uint32_t c = n / 2; c >= 1; c >>= 1) { + } + const auto p = ((lane_id & b) == 0); + for (std::uint32_t c = N / 2; c >= 1; c >>= 1) { #pragma unroll - for (std::uint32_t i = 0; i < n; i++) { - std::uint32_t j = i ^ c; - if (i >= j) continue; - swap_if_needed(ks[i], vs[i], ks[j], vs[j], p); - } + for (std::uint32_t i = 0; i < N; i++) { + std::uint32_t j = i ^ c; + if (i >= j) continue; + swap_if_needed(k[i], v[i], k[j], v[j], p); } } } }; template -struct warp_merge_core_2 { - RAFT_DEVICE_INLINE_FUNCTION void operator()(K* ks, - V* vs, +struct warp_merge_core { + RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[6], + V v[6], + const std::uint32_t range, + const bool asc) + { + constexpr unsigned N = 6; + const auto lane_id = threadIdx.x % warp_size; + + if (range == 1) { + for (std::uint32_t i = 0; i < N; i += 3) { + const auto p = (i == 0); + swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); + swap_if_needed(k[1 + i], v[1 + i], k[2 + i], v[2 + i], p); + swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); + } + const auto p = ((lane_id & 1) == 0); + for (std::uint32_t i = 0; i < 3; i++) { + std::uint32_t j = i + 3; + swap_if_needed(k[i], v[i], k[j], v[j], p); + } + for (std::uint32_t i = 0; i < N; i += 3) { + swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); + swap_if_needed(k[1 + i], v[1 + i], k[2 + i], v[2 + i], p); + swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); + } + return; + } + + const std::uint32_t b = range; + for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { + const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); +#pragma unroll + for (std::uint32_t i = 0; i < N; i++) { + swap_if_needed(k[i], v[i], c, p); + } + } + const auto p = ((lane_id & b) == 0); + for (std::uint32_t i = 0; i < 3; i++) { + std::uint32_t j = i + 3; + swap_if_needed(k[i], v[i], k[j], v[j], p); + } + for (std::uint32_t i = 0; i < N; i += N / 2) { + swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); + swap_if_needed(k[1 + i], v[1 + i], k[2 + i], v[2 + i], p); + swap_if_needed(k[0 + i], v[0 + i], k[1 + i], v[1 + i], p); + } + } +}; + +template +struct warp_merge_core { + RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[3], + V v[3], + const std::uint32_t range, + const bool asc) + { + constexpr unsigned N = 3; + const auto lane_id = threadIdx.x % warp_size; + + if (range == 1) { + const auto p = ((lane_id & 1) == 0); + swap_if_needed(k[0], v[0], k[1], v[1], p); + swap_if_needed(k[1], v[1], k[2], v[2], p); + swap_if_needed(k[0], v[0], k[1], v[1], p); + return; + } + + const std::uint32_t b = range; + for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { + const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); +#pragma unroll + for (std::uint32_t i = 0; i < N; i++) { + swap_if_needed(k[i], v[i], c, p); + } + } + const auto p = ((lane_id & b) == 0); + swap_if_needed(k[0], v[0], k[1], v[1], p); + swap_if_needed(k[1], v[1], k[2], v[2], p); + swap_if_needed(k[0], v[0], k[1], v[1], p); + } +}; + +template +struct warp_merge_core { + RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[2], + V v[2], const std::uint32_t range, const bool asc) { @@ -95,26 +182,27 @@ struct warp_merge_core_2 { if (range == 1) { const auto p = ((lane_id & 1) == 0); - swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); - } else { - const std::uint32_t b = range; - for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { - const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); + swap_if_needed(k[0], v[0], k[1], v[1], p); + return; + } + + const std::uint32_t b = range; + for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { + const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); #pragma unroll - for (std::uint32_t i = 0; i < N; i++) { - swap_if_needed(ks[i], vs[i], c, p); - } + for (std::uint32_t i = 0; i < N; i++) { + swap_if_needed(k[i], v[i], c, p); } - const auto p = ((lane_id & b) == 0); - swap_if_needed(ks[0], vs[0], ks[1], vs[1], p); } + const auto p = ((lane_id & b) == 0); + swap_if_needed(k[0], v[0], k[1], v[1], p); } }; template -struct warp_merge_core_1 { - RAFT_DEVICE_INLINE_FUNCTION void operator()(K* ks, - V* vs, +struct warp_merge_core { + RAFT_DEVICE_INLINE_FUNCTION void operator()(K k[1], + V v[1], const std::uint32_t range, const bool asc) { @@ -122,32 +210,32 @@ struct warp_merge_core_1 { const std::uint32_t b = range; for (std::uint32_t c = b / 2; c >= 1; c >>= 1) { const auto p = static_cast(lane_id & b) == static_cast(lane_id & c); - swap_if_needed(ks[0], vs[0], c, p); + swap_if_needed(k[0], v[0], c, p); } } }; } // namespace detail -template -RAFT_DEVICE_INLINE_FUNCTION void warp_merge( - K* ks, V* vs, unsigned n, unsigned range, const bool asc = true) +template +RAFT_DEVICE_INLINE_FUNCTION void warp_merge(K k[N], V v[N], unsigned range, const bool asc = true) { - if (n == 1) { - detail::warp_merge_core_1{}(ks, vs, range, asc); - } else if (n == 2) { - detail::warp_merge_core_2{}(ks, vs, range, asc); - } else { - detail::warp_merge_core_n{}(ks, vs, n, range, asc); - } + detail::warp_merge_core{}(k, v, range, asc); } -template -RAFT_DEVICE_INLINE_FUNCTION void warp_sort(K* ks, V* vs, unsigned n, const bool asc = true) +template +RAFT_DEVICE_INLINE_FUNCTION void warp_sort(K k[N], V v[N], const bool asc = true) { + if constexpr (N < 8) { +#pragma unroll + for (std::uint32_t range = 1; range <= warp_size; range <<= 1) { + warp_merge(k, v, range, asc); + } + } else { #pragma unroll 1 - for (std::uint32_t range = 1; range <= warp_size; range <<= 1) { - warp_merge(ks, vs, n, range, asc); + for (std::uint32_t range = 1; range <= warp_size; range <<= 1) { + warp_merge(k, v, range, asc); + } } } From 6aae61dc9ee3d106c0d25655e8b2160e92ce3553 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 12 Nov 2025 12:44:12 -0800 Subject: [PATCH 06/25] remove __inline__ from topk_by_radix (to lower register pressure) --- .../neighbors/detail/cagra/topk_by_radix.cuh | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh index da571129a6..9b72f11e8d 100644 --- a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh @@ -18,19 +18,19 @@ struct topk_by_radix_sort_base { template struct topk_by_radix_sort : topk_by_radix_sort_base { - RAFT_DEVICE_INLINE_FUNCTION void operator()(uint32_t max_topk, - uint32_t num_sort_threads, - uint32_t topk, - uint32_t batch_size, - uint32_t len_x, - const uint32_t* _x, - const IdxT* _in_vals, - uint32_t* _y, - IdxT* _out_vals, - uint32_t* work, - uint32_t* _hints, - bool sort, - uint32_t* _smem) + __device__ void operator()(uint32_t max_topk, + uint32_t num_sort_threads, + uint32_t topk, + uint32_t batch_size, + uint32_t len_x, + const uint32_t* _x, + const IdxT* _in_vals, + uint32_t* _y, + IdxT* _out_vals, + uint32_t* work, + uint32_t* _hints, + bool sort, + uint32_t* _smem) { assert(max_topk / num_sort_threads <= 4); assert(blockDim.x >= V / 4); From 8aaf11b1c894e521c5f190cb06c58f83399032a9 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 12 Nov 2025 15:27:42 -0800 Subject: [PATCH 07/25] branch outside topk_by_bitonic_sort --- .../cagra/search_multi_cta_kernel-inl.cuh | 29 ++++++++++--------- 1 file changed, 16 insertions(+), 13 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 879c41487a..4d5fffb972 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -109,23 +109,17 @@ RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parent( } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort(float* distances, // [num_elements] INDEX_T* indices, // [num_elements] - const uint32_t max_elements, const uint32_t num_elements) { const unsigned warp_id = threadIdx.x / raft::warp_size(); if (warp_id > 0) { return; } const unsigned lane_id = threadIdx.x % raft::warp_size(); - assert(max_elements <= 256); - constexpr unsigned MAX_N = - 8; // if MAX_N >> N, we may have negative performance impact, if this is significant, we may - // get memory space from dynamically sized shared memory - const unsigned N = (max_elements + (raft::warp_size() - 1)) / raft::warp_size(); - assert(N <= MAX_N); - float key[MAX_N]; - INDEX_T val[MAX_N]; + constexpr unsigned N = (MAX_ELEMENTS + (raft::warp_size() - 1)) / raft::warp_size(); + float key[N]; + INDEX_T val[N]; for (unsigned i = 0; i < N; i++) { unsigned j = lane_id + (raft::warp_size() * i); if (j < num_elements) { @@ -137,7 +131,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort(float* distances, // [num } } /* Warp Sort */ - bitonic::warp_sort(key, val, N); + bitonic::warp_sort(key, val); /* Store sorted results */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * lane_id) + i; @@ -272,8 +266,17 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( _CLK_START(); if (threadIdx.x < 32) { // [1st warp] Topk with bitonic sort - topk_by_bitonic_sort( - result_distances_buffer, result_indices_buffer, max_elements, result_buffer_size_32); + if (max_elements <= 64) { + topk_by_bitonic_sort<64, INDEX_T>( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } else if (max_elements <= 128) { + topk_by_bitonic_sort<128, INDEX_T>( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } else { + assert(max_elements <= 256); + topk_by_bitonic_sort<256, INDEX_T>( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } } __syncthreads(); _CLK_REC(clk_topk); From 07351299928763e986eb177b0b65862555c1b82d Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 12 Nov 2025 15:46:11 -0800 Subject: [PATCH 08/25] use shared memory when I need to create large stack arrays for bitonic sort key, value pairs --- .../detail/cagra/search_single_cta.cuh | 46 +++++++++++-------- 1 file changed, 27 insertions(+), 19 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh index 9d59a74df2..986bab0ce0 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh @@ -129,13 +129,34 @@ struct search sizeof(std::uint32_t) * topk_ws_size + sizeof(std::uint32_t); std::uint32_t additional_smem_size = 0; - if (num_itopk_candidates > 256) { - // Tentatively calculate the required share memory size when radix - // sort based topk is used, assuming the block size is the maximum. + if (num_itopk_candidates <= 256) { // bitonic sort + // Tentatively calculate the required shared memory size when bitonic sort based topk is used, + // assuming the block size is the maximum. + uint32_t max_elements{0}; + if (itopk_size > 128) { // use shared memory instead of stack variables (which may or may not + // be placed in registers) for bitonic sort + if (itopk_size <= 256) { + max_elements = 256; + } else { + static_assert(min_block_size >= + 64); // multi_warps = true if (threadDim.x >= 64) && (itopk_size > 256) + max_elements = 512 / 2 /* multi_warps = true */; + } + } + additional_smem_size += max_block_size * + ((max_elements + (raft::warp_size() - 1)) / raft::warp_size()) * + (sizeof(float) + sizeof(INDEX_T)); + } else { // radix sort + // Tentatively calculate the required shared memory size when radix sort based topk is used, + // assuming the block size is the maximum. if (itopk_size <= 256) { - additional_smem_size += topk_by_radix_sort::smem_size(256) * sizeof(std::uint32_t); + constexpr unsigned MAX_ITOPK = 256; + additional_smem_size += + topk_by_radix_sort::smem_size(MAX_ITOPK) * sizeof(std::uint32_t); } else { - additional_smem_size += topk_by_radix_sort::smem_size(512) * sizeof(std::uint32_t); + constexpr unsigned MAX_ITOPK = 512; + additional_smem_size += + topk_by_radix_sort::smem_size(MAX_ITOPK) * sizeof(std::uint32_t); } } @@ -152,7 +173,7 @@ struct search if (block_size == 0) { block_size = min_block_size; - if (num_itopk_candidates > 256) { + if (num_itopk_candidates > 256) { // radix sort // radix-based topk is used. block_size = min_block_size_radix; @@ -190,19 +211,6 @@ struct search max_block_size); thread_block_size = block_size; - if (num_itopk_candidates <= 256) { - RAFT_LOG_DEBUG("# bitonic-sort based topk routine is used"); - } else { - RAFT_LOG_DEBUG("# radix-sort based topk routine is used"); - smem_size = base_smem_size; - if (itopk_size <= 256) { - constexpr unsigned MAX_ITOPK = 256; - smem_size += topk_by_radix_sort::smem_size(MAX_ITOPK) * sizeof(std::uint32_t); - } else { - constexpr unsigned MAX_ITOPK = 512; - smem_size += topk_by_radix_sort::smem_size(MAX_ITOPK) * sizeof(std::uint32_t); - } - } RAFT_LOG_DEBUG("# smem_size: %u", smem_size); hashmap_size = 0; if (small_hash_bitlen == 0 && !this->persistent) { From 3c9ee5f6a958158ab73278b9a789da811cb6aa9e Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 12 Nov 2025 16:36:39 -0800 Subject: [PATCH 09/25] branch before calling topky_by_bitonic_sort_and_merge --- .../cagra/search_single_cta_kernel-inl.cuh | 233 +++++++++++------- 1 file changed, 139 insertions(+), 94 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index ed6f18daa6..4ab640e63b 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -55,6 +55,8 @@ namespace cuvs::neighbors::cagra::detail { namespace single_cta_search { +struct empty_t {}; + // #define _CLK_BREAKDOWN template @@ -99,27 +101,23 @@ RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parents(std::uint32_t* const termin if (threadIdx.x == 0 && (num_new_parents == 0)) { *terminate_flag = 1; } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] - const std::uint32_t max_candidates, const std::uint32_t num_candidates, const std::uint32_t num_itopk, - bool multi_warps = false) + uint32_t* _smem) { const unsigned lane_id = threadIdx.x % raft::warp_size(); const unsigned warp_id = threadIdx.x / raft::warp_size(); - assert(max_candidates <= 256); - constexpr unsigned MAX_N = - 8; // if MAX_N >> N, we may have negative performance impact, if this is significant, we may - // get memory space from dynamically sized shared memory. - float key[MAX_N]; - IdxT val[MAX_N]; - if (multi_warps == false) { + static_assert(MAX_CANDIDATES <= 256); + if constexpr (!MULTI_WARPS) { + static_assert(MAX_CANDIDATES <= 128); if (warp_id > 0) { return; } - const unsigned N = (max_candidates + (raft::warp_size() - 1)) / raft::warp_size(); - assert(N <= MAX_N); + constexpr unsigned N = (MAX_CANDIDATES + (raft::warp_size() - 1)) / raft::warp_size(); + float key[N]; + IdxT val[N]; /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { unsigned j = lane_id + (raft::warp_size() * i); @@ -132,7 +130,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } /* Sort */ - bitonic::warp_sort(key, val, N); + bitonic::warp_sort(key, val); /* Reg -> Temp_itopk */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * lane_id) + i; @@ -143,9 +141,11 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } else { // Use two warps (64 threads) - const unsigned max_candidates_per_warp = (max_candidates + 1) / 2; - const unsigned N = (max_candidates_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); - assert(N <= MAX_N); + constexpr unsigned max_candidates_per_warp = (MAX_CANDIDATES + 1) / 2; + static_assert(max_candidates_per_warp <= 128); + constexpr unsigned N = (max_candidates_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); + float key[N]; + IdxT val[N]; if (warp_id < 2) { /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { @@ -160,7 +160,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } /* Sort */ - bitonic::warp_sort(key, val, N); + bitonic::warp_sort(key, val); /* Reg -> Temp_candidates */ for (unsigned i = 0; i < N; i++) { unsigned jl = (N * lane_id) + i; @@ -180,7 +180,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( unsigned jl = (N * lane_id) + i; unsigned kl = max_candidates_per_warp - 1 - jl; unsigned j = jl + (max_candidates_per_warp * warp_id); - unsigned k = max_candidates - 1 - j; + unsigned k = MAX_CANDIDATES - 1 - j; if (j >= num_candidates || k >= num_candidates || kl >= num_itopk) continue; float temp_key = candidate_distances[device::swizzling(k)]; if (key[i] == temp_key) continue; @@ -193,7 +193,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( if (num_warps_used > 1) { __syncthreads(); } if (warp_id < num_warps_used) { /* Merge */ - bitonic::warp_merge(key, val, N, raft::warp_size()); + bitonic::warp_merge(key, val, raft::warp_size()); /* Reg -> Temp_itopk */ for (unsigned i = 0; i < N; i++) { unsigned jl = (N * lane_id) + i; @@ -208,32 +208,36 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] - const std::uint32_t max_itopk, const std::uint32_t num_itopk, float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, std::uint32_t* work_buf, const bool first, - bool multi_warps = false) + uint32_t* _smem) { const unsigned lane_id = threadIdx.x % raft::warp_size(); const unsigned warp_id = threadIdx.x / raft::warp_size(); - assert(max_itopk <= 512); - constexpr unsigned MAX_N = - 16; // if MAX_N >> N, we may have negative performance impact, if this is significant, we may - // get memory space from dynamically sized shared memory. - float key[MAX_N]; - IdxT val[MAX_N]; - if (multi_warps == false) { + static_assert(MAX_ITOPK <= 512); + if constexpr (!MULTI_WARPS) { if (warp_id > 0) { return; } - const unsigned N = (max_itopk + (raft::warp_size() - 1)) / raft::warp_size(); - assert(N <= MAX_N); + constexpr unsigned N = (MAX_ITOPK + (raft::warp_size() - 1)) / raft::warp_size(); + float* key{}; + IdxT* val{}; + std::conditional_t<(MAX_ITOPK <= 128), float[N], empty_t> stack_key{}; + std::conditional_t<(MAX_ITOPK <= 128), IdxT[N], empty_t> stack_val{}; + if constexpr (MAX_ITOPK <= 128) { + key = stack_key; + val = stack_val; + } else { + key = reinterpret_cast(_smem); + val = reinterpret_cast(_smem + N * sizeof(float)); + } if (first) { /* Load itopk results */ for (unsigned i = 0; i < N; i++) { @@ -247,7 +251,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Sort */ - bitonic::warp_sort(key, val, N); + bitonic::warp_sort(key, val); } else { /* Load itopk results */ for (unsigned i = 0; i < N; i++) { @@ -264,7 +268,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( /* Merge candidates */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * lane_id) + i; // [0:max_itopk-1] - unsigned k = max_itopk - 1 - j; + unsigned k = MAX_ITOPK - 1 - j; if (k >= num_itopk || k >= num_candidates) continue; float candidate_key = candidate_distances[device::swizzling(k)]; if (key[i] > candidate_key) { @@ -273,7 +277,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Merge */ - bitonic::warp_merge(key, val, N, raft::warp_size()); + bitonic::warp_merge(key, val, raft::warp_size()); /* Store new itopk results */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * lane_id) + i; @@ -284,9 +288,19 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } else { // Use two warps (64 threads) or more - const unsigned max_itopk_per_warp = (max_itopk + 1) / 2; - const unsigned N = (max_itopk_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); - assert(N <= MAX_N); + constexpr unsigned max_itopk_per_warp = (MAX_ITOPK + 1) / 2; + constexpr unsigned N = (max_itopk_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); + float* key{}; + IdxT* val{}; + std::conditional_t<(max_itopk_per_warp <= 128), float[N], empty_t> stack_key{}; + std::conditional_t<(max_itopk_per_warp <= 128), IdxT[N], empty_t> stack_val{}; + if constexpr (max_itopk_per_warp <= 128) { + key = stack_key; + val = stack_val; + } else { + key = reinterpret_cast(_smem); + val = reinterpret_cast(_smem + N * sizeof(float)); + } if (first) { /* Load itop results (not sorted) */ if (warp_id < 2) { @@ -301,7 +315,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Sort */ - bitonic::warp_sort(key, val, N); + bitonic::warp_sort(key, val); /* Store intermedidate results */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * threadIdx.x) + i; @@ -315,7 +329,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( /* Load intermedidate results */ for (unsigned i = 0; i < N; i++) { unsigned j = (N * threadIdx.x) + i; - unsigned k = max_itopk - 1 - j; + unsigned k = MAX_ITOPK - 1 - j; if (k >= num_itopk) continue; float temp_key = itopk_distances[device::swizzling(k)]; if (key[i] == temp_key) continue; @@ -325,7 +339,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Merge */ - bitonic::warp_merge(key, val, N, raft::warp_size()); + bitonic::warp_merge(key, val, raft::warp_size()); } __syncthreads(); /* Store itopk results (sorted) */ @@ -394,8 +408,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( unsigned j = (N * lane_id) + i; if (j < turning_point) { k = j + (num_itopk_div2 * warp_id); - } else if (j >= (max_itopk / 2 - num_itopk_div2)) { - j -= (max_itopk / 2 - num_itopk_div2); + } else if (j >= (MAX_ITOPK / 2 - num_itopk_div2)) { + j -= (MAX_ITOPK / 2 - num_itopk_div2); if ((turning_point <= j) && (j < num_itopk_div2)) { k = j + (num_itopk_div2 * warp_id); } } if (k < num_itopk) { @@ -407,7 +421,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } /* Warp Merge */ - bitonic::warp_merge(key, val, N, raft::warp_size()); + bitonic::warp_merge(key, val, raft::warp_size()); /* Store new itopk results */ for (unsigned i = 0; i < N; i++) { const unsigned j = (N * lane_id) + i; @@ -422,40 +436,86 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } template -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( - float* itopk_distances, // [num_itopk] - IdxT* itopk_indices, // [num_itopk] - const std::uint32_t max_itopk, - const std::uint32_t num_itopk, - float* candidate_distances, // [num_candidates] - IdxT* candidate_indices, // [num_candidates] - const std::uint32_t max_candidates, - const std::uint32_t num_candidates, - std::uint32_t* work_buf, - const bool first, - const bool multi_warps_1, - const bool multi_warps_2) +__device__ void topk_by_bitonic_sort_and_merge(float* itopk_distances, // [num_itopk] + IdxT* itopk_indices, // [num_itopk] + const std::uint32_t max_itopk, + const std::uint32_t num_itopk, + float* candidate_distances, // [num_candidates] + IdxT* candidate_indices, // [num_candidates] + const std::uint32_t max_candidates, + const std::uint32_t num_candidates, + std::uint32_t* work_buf, + const bool first, + uint32_t* _smem) { // The results in candidate_distances/indices are sorted by bitonic sort. - topk_by_bitonic_sort_and_full(candidate_distances, - candidate_indices, - max_candidates, - num_candidates, - num_itopk, - multi_warps_1); + assert(blockDim.x >= 64); + if (max_candidates <= 64) { + constexpr bool MULTI_WARPS_1 = false; + topk_by_bitonic_sort_and_full<64, MULTI_WARPS_1, IdxT>(candidate_distances, + candidate_indices, + num_candidates, + num_itopk, + static_cast(nullptr)); + } else if (max_candidates <= 128) { + constexpr bool MULTI_WARPS_1 = false; + topk_by_bitonic_sort_and_full<128, MULTI_WARPS_1, IdxT>( + candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); + } else { + constexpr bool MULTI_WARPS_1 = true; + assert(max_candidates <= 256); + topk_by_bitonic_sort_and_full<256, MULTI_WARPS_1, IdxT>( + candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); + } // The results sorted above are merged with the internal intermediate top-k // results so far using bitonic merge. - topk_by_bitonic_sort_and_merge(itopk_distances, - itopk_indices, - max_itopk, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - multi_warps_2); + if (max_itopk <= 64) { + constexpr bool MULTI_WARPS_2 = false; + topk_by_bitonic_sort_and_merge<64, MULTI_WARPS_2, IdxT>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + static_cast(nullptr)); + } else if (max_itopk <= 128) { + constexpr bool MULTI_WARPS_2 = false; + topk_by_bitonic_sort_and_merge<128, MULTI_WARPS_2, IdxT>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); + } else if (max_itopk <= 256) { + constexpr bool MULTI_WARPS_2 = false; + topk_by_bitonic_sort_and_merge<256, MULTI_WARPS_2, IdxT>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); + } else { + assert(max_itopk <= 512); + constexpr bool MULTI_WARPS_2 = true; + topk_by_bitonic_sort_and_merge<512, MULTI_WARPS_2, IdxT>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); + } } // This function move the invalid index element to the end of the itopk list. @@ -681,8 +741,9 @@ __device__ void search_core( // topk_by_bitonic_sort_and_merge() consists of two operations: // if max_candidates is greater than 128, the first operation uses two warps; // if max_itopk is greater than 256, the second operation used two warps. - const bool multi_warps_1 = ((blockDim.x >= 64) && (max_candidates > 128)) ? true : false; - const bool multi_warps_2 = ((blockDim.x >= 64) && (max_itopk > 256)) ? true : false; + assert(blockDim.x >= 64); + const bool multi_warps_1 = (max_candidates > 128) ? true : false; + const bool multi_warps_2 = (max_itopk > 256) ? true : false; // reset small-hash table. if ((iter + 1) % small_hash_reset_interval == 0) { @@ -733,8 +794,7 @@ __device__ void search_core( search_width * graph_degree, topk_ws, (iter == 0), - multi_warps_1, - multi_warps_2); + smem_work_ptr); __syncthreads(); _CLK_REC(clk_topk); } else { @@ -892,8 +952,6 @@ __device__ void search_core( // candidate list. if (top_k > internal_topk || result_indices_buffer[top_k - 1] == invalid_index) { __syncthreads(); - const bool multi_warps_1 = ((blockDim.x >= 64) && (max_candidates > 128)) ? true : false; - const bool multi_warps_2 = ((blockDim.x >= 64) && (max_itopk > 256)) ? true : false; topk_by_bitonic_sort_and_merge(result_distances_buffer, result_indices_buffer, max_itopk, @@ -904,8 +962,7 @@ __device__ void search_core( search_width * graph_degree, topk_ws, (iter == 0), - multi_warps_1, - multi_warps_2); + smem_work_ptr); } __syncthreads(); } @@ -1231,25 +1288,13 @@ struct search_kernel_config { using kernel_t = decltype(dispatch_kernel); - template - static auto choose_search_kernel(unsigned itopk_size) -> kernel_t - { - if (itopk_size <= 512) { - return dispatch_kernel; - } - THROW("No kernel for parameters itopk_size %u", itopk_size); - } - static auto choose_itopk_and_mx_candidates(unsigned itopk_size, unsigned num_itopk_candidates, unsigned block_size) -> kernel_t { if (num_itopk_candidates <= 256) { - return choose_search_kernel<1>(itopk_size); + assert(itopk_size <= 512); + return dispatch_kernel; } else { // Radix-based topk is used return dispatch_kernel; From 68b191bcc47d701c811532f3de40f6c66c88495f Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 12 Nov 2025 22:43:44 -0800 Subject: [PATCH 10/25] undo changes in topk_cta_11_core (branch in the caller site) --- .../neighbors/detail/cagra/topk_by_radix.cuh | 98 +++++++++++--- .../detail/cagra/topk_for_cagra/topk.cu | 124 +++++++++--------- .../detail/cagra/topk_for_cagra/topk_core.cuh | 104 +++++++-------- 3 files changed, 187 insertions(+), 139 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh index 9b72f11e8d..a86c011f18 100644 --- a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh @@ -19,9 +19,7 @@ struct topk_by_radix_sort_base { template struct topk_by_radix_sort : topk_by_radix_sort_base { __device__ void operator()(uint32_t max_topk, - uint32_t num_sort_threads, uint32_t topk, - uint32_t batch_size, uint32_t len_x, const uint32_t* _x, const IdxT* _in_vals, @@ -32,24 +30,88 @@ struct topk_by_radix_sort : topk_by_radix_sort_base { bool sort, uint32_t* _smem) { - assert(max_topk / num_sort_threads <= 4); assert(blockDim.x >= V / 4); std::uint8_t* const state = reinterpret_cast(work); - topk_cta_11_core(max_topk, - num_sort_threads, - topk, - len_x, - _x, - _in_vals, - _y, - _out_vals, - state, - _hints, - sort, - _smem); + if (max_topk <= 64) { + topk_cta_11_core(topk, + len_x, + _x, + _in_vals, + _y, + _out_vals, + state, + _hints, + sort, + _smem); + } + else if (max_topk <= 128) { + topk_cta_11_core(topk, + len_x, + _x, + _in_vals, + _y, + _out_vals, + state, + _hints, + sort, + _smem); + } + else if (max_topk <= 256) { + topk_cta_11_core(topk, + len_x, + _x, + _in_vals, + _y, + _out_vals, + state, + _hints, + sort, + _smem); + } + else if (max_topk <= 512) { + topk_cta_11_core(topk, + len_x, + _x, + _in_vals, + _y, + _out_vals, + state, + _hints, + sort, + _smem); + } + else { + topk_cta_11_core(topk, + len_x, + _x, + _in_vals, + _y, + _out_vals, + state, + _hints, + sort, + _smem); + } } }; diff --git a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk.cu b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk.cu index 41de68cede..b8c343414e 100644 --- a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk.cu +++ b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk.cu @@ -60,8 +60,6 @@ void _cuann_find_topk(uint32_t topK, dim3 blocks(sizeBatch, 1, 1); void (*cta_kernel)(uint32_t, - uint32_t, - uint32_t, uint32_t, uint32_t, const uint32_t*, @@ -76,73 +74,69 @@ void _cuann_find_topk(uint32_t topK, uint32_t*, bool) = nullptr; - int _vecLen = _get_vecLen(ldIK, 2); - constexpr int maxTopkPerThread = 4; - + // V:vecLen, K:maxTopk, T:numSortThreads +#define SET_KERNEL_VKT(V, K, T, ValT) \ + do { \ + assert(numThreads >= T); \ + assert((K % T) == 0); \ + assert((K / T) <= 4); \ + cta_kernel = kern_topk_cta_11; \ + } while (0) + + // V: vecLen +#define SET_KERNEL_V(V, ValT) \ + do { \ + if (topK <= 32) { \ + SET_KERNEL_VKT(V, 32, 32, ValT); \ + } else if (topK <= 64) { \ + SET_KERNEL_VKT(V, 64, 32, ValT); \ + } else if (topK <= 96) { \ + SET_KERNEL_VKT(V, 96, 32, ValT); \ + } else if (topK <= 128) { \ + SET_KERNEL_VKT(V, 128, 32, ValT); \ + } else if (topK <= 192) { \ + SET_KERNEL_VKT(V, 192, 64, ValT); \ + } else if (topK <= 256) { \ + SET_KERNEL_VKT(V, 256, 64, ValT); \ + } else if (topK <= 384) { \ + SET_KERNEL_VKT(V, 384, 128, ValT); \ + } else if (topK <= 512) { \ + SET_KERNEL_VKT(V, 512, 128, ValT); \ + } else if (topK <= 768) { \ + SET_KERNEL_VKT(V, 768, 256, ValT); \ + } else if (topK <= 1024) { \ + SET_KERNEL_VKT(V, 1024, 256, ValT); \ + } \ + /* else if (topK <= 1536) { SET_KERNEL_VKT(V, 1536, 512); } */ \ + /* else if (topK <= 2048) { SET_KERNEL_VKT(V, 2048, 512); } */ \ + /* else if (topK <= 3072) { SET_KERNEL_VKT(V, 3072, 1024); } */ \ + /* else if (topK <= 4096) { SET_KERNEL_VKT(V, 4096, 1024); } */ \ + else { \ + RAFT_FAIL("topk must be lower than or equal to 1024"); \ + } \ + } while (0) + + int _vecLen = _get_vecLen(ldIK, 2); if (_vecLen == 2) { - cta_kernel = kern_topk_cta_11; + SET_KERNEL_V(2, ValT); } else if (_vecLen == 1) { - cta_kernel = kern_topk_cta_11; - } - - int max_topk{}; - int num_sort_threads{}; - if (topK <= 32) { - max_topk = 32; - num_sort_threads = 32; - } else if (topK <= 64) { - max_topk = 64; - num_sort_threads = 32; - } else if (topK <= 96) { - max_topk = 96; - num_sort_threads = 32; - } else if (topK <= 128) { - max_topk = 128; - num_sort_threads = 32; - } else if (topK <= 192) { - max_topk = 192; - num_sort_threads = 64; - } else if (topK <= 256) { - max_topk = 256; - num_sort_threads = 64; - } else if (topK <= 384) { - max_topk = 384; - num_sort_threads = 128; - } else if (topK <= 512) { - max_topk = 512; - num_sort_threads = 128; - } else if (topK <= 768) { - max_topk = 768; - num_sort_threads = 256; - } else if (topK <= 1024) { - max_topk = 1024; - num_sort_threads = 256; - } else { - RAFT_FAIL("topK must be lower than or equal to 1024"); + SET_KERNEL_V(1, ValT); } - assert(max_topk % num_sort_threads == 0); - assert(max_topk / num_sort_threads <= maxTopkPerThread); - - const size_t smem_len = 2 * max_topk + 2048 + 8; - assert(max_topk * (1 + utils::size_of() / utils::size_of()) <= smem_len); - const size_t smem_size = smem_len * sizeof(uint32_t); - cta_kernel<<>>(max_topk, - num_sort_threads, - topK, - sizeBatch, - numElements, - (const uint32_t*)inputKeys, - ldIK, - inputVals, - ldIV, - (uint32_t*)outputKeys, - ldOK, - outputVals, - ldOV, - state, - hints, - sort); + cta_kernel<<>>(topK, + sizeBatch, + numElements, + (const uint32_t*)inputKeys, + ldIK, + inputVals, + ldIV, + (uint32_t*)outputKeys, + ldOK, + outputVals, + ldOV, + state, + hints, + sort); return; } diff --git a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh index 6e07587aa3..1e9bd86389 100644 --- a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh @@ -599,11 +599,9 @@ RAFT_INLINE_FUNCTION constexpr uint32_t get_state_size(uint32_t len_x) return 0; } -// max_topk / num_sort_threads should never exceed maxTopkPerThread -template -RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, - uint32_t num_sort_threads, - uint32_t topk, +// +template +RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, uint32_t len_x, const uint32_t* _x, // [size_batch, ld_x,] const ValT* _in_vals, // [size_batch, ld_iv,] @@ -615,9 +613,9 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, uint32_t* _smem) { uint32_t* const smem_out_vals = _smem; - uint32_t* const hist = &(_smem[2 * max_topk]); - uint32_t* const best_index = &(_smem[2 * max_topk + 2048]); - uint32_t* const best_csum = &(_smem[2 * max_topk + 2048 + 3]); + uint32_t* const hist = &(_smem[2 * maxTopk]); + uint32_t* const best_index = &(_smem[2 * maxTopk + 2048]); + uint32_t* const best_csum = &(_smem[2 * maxTopk + 2048 + 3]); const uint32_t num_threads = blockDim.x; const uint32_t thread_id = threadIdx.x; @@ -633,11 +631,11 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, const uint32_t hint = (_hint == NULL ? ~0u : *_hint); // Initialize shared memory - for (int i = 2 * max_topk + thread_id; i < 2 * max_topk + 2048 + 8; i += num_threads) { + for (int i = 2 * maxTopk + thread_id; i < 2 * maxTopk + 2048 + 8; i += num_threads) { _smem[i] = 0; } - uint32_t* const output_count = &(_smem[2 * max_topk + 2048 + 6]); - uint32_t* const output_count_eq = &(_smem[2 * max_topk + 2048 + 7]); + uint32_t* const output_count = &(_smem[2 * maxTopk + 2048 + 6]); + uint32_t* const output_count_eq = &(_smem[2 * maxTopk + 2048 + 7]); uint32_t threshold = 0; uint32_t nx_below_threshold = 0; __syncthreads(); @@ -645,7 +643,6 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, // // Search for the maximum threshold that satisfies "(x < threshold).sum() <= topk". // -#pragma unroll for (int j = 0; j < 3; j += 1) { uint32_t num_bins; uint32_t shift; @@ -723,18 +720,14 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, return; } - const int num_topk_per_thread = max_topk / num_sort_threads; - static_assert( - maxTopkPerThread <= - 4); // we will end-up wasting too much stack memory if maxTopkPerThread >> num_topk_per_thread - assert(num_topk_per_thread <= maxTokenPerThread); - float my_keys[maxTopkPerThread]; - ValT my_vals[maxTopkPerThread]; + constexpr int numTopkPerThread = maxTopk / numSortThreads; + float my_keys[numTopkPerThread]; + ValT my_vals[numTopkPerThread]; // Read keys and values to registers - if (thread_id < num_sort_threads) { - for (int i = 0; i < num_topk_per_thread; i++) { - const int k = thread_id + (num_sort_threads * i); + if (thread_id < numSortThreads) { + for (int i = 0; i < numTopkPerThread; i++) { + const int k = thread_id + (numSortThreads * i); if (k < topk) { const int j = smem_out_vals[k]; my_keys[i] = ((float*)x)[j]; @@ -753,21 +746,21 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, uint32_t mask = 1; // Sorting by thread - if (thread_id < num_sort_threads) { + if (thread_id < numSortThreads) { const bool ascending = ((thread_id & mask) == 0); - if (num_topk_per_thread == 3) { + if constexpr (numTopkPerThread == 3) { swap_if_needed(my_keys[0], my_keys[1], my_vals[0], my_vals[1], ascending); swap_if_needed(my_keys[0], my_keys[2], my_vals[0], my_vals[2], ascending); swap_if_needed(my_keys[1], my_keys[2], my_vals[1], my_vals[2], ascending); } else { - for (int j = 0; j < num_topk_per_thread / 2; j += 1) { + for (int j = 0; j < numTopkPerThread / 2; j += 1) { #pragma unroll - for (int i = 0; i < num_topk_per_thread; i += 2) { + for (int i = 0; i < numTopkPerThread; i += 2) { swap_if_needed( my_keys[i], my_keys[i + 1], my_vals[i], my_vals[i + 1], ascending); } #pragma unroll - for (int i = 1; i < num_topk_per_thread - 1; i += 2) { + for (int i = 1; i < numTopkPerThread - 1; i += 2) { swap_if_needed( my_keys[i], my_keys[i + 1], my_vals[i], my_vals[i + 1], ascending); } @@ -776,38 +769,38 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, } // Bitonic Sorting - while (mask < num_sort_threads) { + while (mask < numSortThreads) { uint32_t next_mask = mask << 1; for (uint32_t curr_mask = mask; curr_mask > 0; curr_mask >>= 1) { const bool ascending = ((thread_id & curr_mask) == 0) == ((thread_id & next_mask) == 0); if (curr_mask >= 32) { // inter warp - ValT* const smem_vals = reinterpret_cast(_smem); // [max_topk] - float* const smem_keys = reinterpret_cast( - smem_vals + max_topk); // [num_topk_per_thread, num_sort_threads] + ValT* const smem_vals = reinterpret_cast(_smem); // [maxTopk] + float* const smem_keys = + reinterpret_cast(smem_vals + maxTopk); // [numTopkPerThread, numSortThreads] __syncthreads(); - if (thread_id < num_sort_threads) { + if (thread_id < numSortThreads) { #pragma unroll - for (int i = 0; i < num_topk_per_thread; i++) { - smem_keys[thread_id + (num_sort_threads * i)] = my_keys[i]; - smem_vals[thread_id + (num_sort_threads * i)] = my_vals[i]; + for (int i = 0; i < numTopkPerThread; i++) { + smem_keys[thread_id + (numSortThreads * i)] = my_keys[i]; + smem_vals[thread_id + (numSortThreads * i)] = my_vals[i]; } } __syncthreads(); - if (thread_id < num_sort_threads) { + if (thread_id < numSortThreads) { #pragma unroll - for (int i = 0; i < num_topk_per_thread; i++) { - float opp_key = smem_keys[(thread_id ^ curr_mask) + (num_sort_threads * i)]; - ValT opp_val = smem_vals[(thread_id ^ curr_mask) + (num_sort_threads * i)]; + for (int i = 0; i < numTopkPerThread; i++) { + float opp_key = smem_keys[(thread_id ^ curr_mask) + (numSortThreads * i)]; + ValT opp_val = smem_vals[(thread_id ^ curr_mask) + (numSortThreads * i)]; swap_if_needed(my_keys[i], opp_key, my_vals[i], opp_val, ascending); } } } else { // intra warp - if (thread_id < num_sort_threads) { + if (thread_id < numSortThreads) { #pragma unroll - for (int i = 0; i < num_topk_per_thread; i++) { + for (int i = 0; i < numTopkPerThread; i++) { float opp_key = __shfl_xor_sync(0xffffffff, my_keys[i], curr_mask); ValT opp_val = __shfl_xor_sync(0xffffffff, my_vals[i], curr_mask); swap_if_needed(my_keys[i], opp_key, my_vals[i], opp_val, ascending); @@ -816,17 +809,17 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, } } - if (thread_id < num_sort_threads) { + if (thread_id < numSortThreads) { const bool ascending = ((thread_id & next_mask) == 0); - if (num_topk_per_thread == 3) { + if constexpr (numTopkPerThread == 3) { swap_if_needed(my_keys[0], my_keys[1], my_vals[0], my_vals[1], ascending); swap_if_needed(my_keys[0], my_keys[2], my_vals[0], my_vals[2], ascending); swap_if_needed(my_keys[1], my_keys[2], my_vals[1], my_vals[2], ascending); } else { #pragma unroll - for (uint32_t curr_mask = num_topk_per_thread / 2; curr_mask > 0; curr_mask >>= 1) { + for (uint32_t curr_mask = numTopkPerThread / 2; curr_mask > 0; curr_mask >>= 1) { #pragma unroll - for (int i = 0; i < num_topk_per_thread; i++) { + for (int i = 0; i < numTopkPerThread; i++) { const int j = i ^ curr_mask; if (i > j) continue; swap_if_needed(my_keys[i], my_keys[j], my_vals[i], my_vals[j], ascending); @@ -838,9 +831,9 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t max_topk, } // Write sorted keys and values - if (thread_id < num_sort_threads) { - for (int i = 0; i < num_topk_per_thread; i++) { - const int k = i + (num_topk_per_thread * thread_id); + if (thread_id < numSortThreads) { + for (int i = 0; i < numTopkPerThread; i++) { + const int k = i + (numTopkPerThread * thread_id); if (k < topk) { if (y) { y[k] = reinterpret_cast(my_keys)[i]; } if (out_vals) { out_vals[k] = my_vals[i]; } @@ -868,11 +861,9 @@ int _get_vecLen(uint32_t maxSamples, int maxVecLen = MAX_VEC_LENGTH) } } // unnamed namespace -template +template __launch_bounds__(1024, 1) RAFT_KERNEL - kern_topk_cta_11(uint32_t max_topk, - uint32_t num_sort_threads, - uint32_t topk, + kern_topk_cta_11(uint32_t topk, uint32_t size_batch, uint32_t len_x, const uint32_t* _x, // [size_batch, ld_x,] @@ -890,11 +881,12 @@ __launch_bounds__(1024, 1) RAFT_KERNEL const uint32_t i_batch = blockIdx.x; if (i_batch >= size_batch) return; - extern __shared__ uint32_t _smem[]; + constexpr uint32_t smem_len = 2 * maxTopk + 2048 + 8; + static_assert(maxTopk * (1 + utils::size_of() / utils::size_of()) <= smem_len, + "maxTopk * sizeof(ValT) must be smaller or equal to 8192 byte"); + __shared__ uint32_t _smem[smem_len]; - topk_cta_11_core( - max_topk, - num_sort_threads, + topk_cta_11_core( topk, len_x, (_x == NULL ? NULL : _x + i_batch * ld_x), From 89d9698ed71089805e36a5bbc55bf058454df47b Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Thu, 13 Nov 2025 15:07:58 -0800 Subject: [PATCH 11/25] fix build error --- cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 4ab640e63b..306e6736a5 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -801,9 +801,7 @@ __device__ void search_core( _CLK_START(); // topk with radix block sort topk_by_radix_sort{}(max_itopk, - (max_itopk <= 64) ? 32 : max_itopk / 4, internal_topk, - gridDim.x, result_buffer_size, reinterpret_cast(result_distances_buffer), result_indices_buffer, From a9bfab7a0dc33e7f4662ce8f5ac63fbaff6f2987 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Thu, 13 Nov 2025 17:20:44 -0800 Subject: [PATCH 12/25] update max_itopk setting in single-CTA radix sort based search to match the original code --- .../cagra/search_single_cta_kernel-inl.cuh | 38 ++++--- .../neighbors/detail/cagra/topk_by_radix.cuh | 98 +++---------------- 2 files changed, 40 insertions(+), 96 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 306e6736a5..788ae3336e 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -1290,8 +1290,8 @@ struct search_kernel_config { unsigned num_itopk_candidates, unsigned block_size) -> kernel_t { + assert(itopk_size <= 512); if (num_itopk_candidates <= 256) { - assert(itopk_size <= 512); return dispatch_kernel; } else { // Radix-based topk is used @@ -2141,19 +2141,6 @@ void select_and_run( const SourceIndexT* source_indices_ptr = source_indices.has_value() ? source_indices->data_handle() : nullptr; - uint32_t max_itopk{}; - if (ps.itopk_size <= 64) { - max_itopk = 64; - } else if (ps.itopk_size <= 128) { - max_itopk = 128; - } else if (ps.itopk_size <= 256) { - max_itopk = 256; - } else if (ps.itopk_size <= 512) { - max_itopk = 512; - } else { - THROW("No kernel for parameter itopk_size %u", ps.itopk_size); - } - uint32_t max_candidates{}; if (num_itopk_candidates <= 64) { max_candidates = 64; @@ -2166,6 +2153,29 @@ void select_and_run( 32; // irrelevant, radix based topk is used (see choose_itopk_and_max_candidates) } + uint32_t max_itopk{}; + if (num_itopk_candidates <= 256) { // bitonic sort + if (ps.itopk_size <= 64) { + max_itopk = 64; + } else if (ps.itopk_size <= 128) { + max_itopk = 128; + } else if (ps.itopk_size <= 256) { + max_itopk = 256; + } else if (ps.itopk_size <= 512) { + max_itopk = 512; + } else { + THROW("No kernel for parameter itopk_size %u", ps.itopk_size); + } + } else { // radix sort + if (ps.itopk_size <= 256) { + max_itopk = 256; + } else if (ps.itopk_size <= 512) { + max_itopk = 512; + } else { + THROW("No kernel for parameter itopk_size %u", ps.itopk_size); + } + } + if (ps.persistent) { using runner_type = persistent_runner_t; diff --git a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh index a86c011f18..0d1e679a01 100644 --- a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh @@ -18,99 +18,33 @@ struct topk_by_radix_sort_base { template struct topk_by_radix_sort : topk_by_radix_sort_base { - __device__ void operator()(uint32_t max_topk, - uint32_t topk, - uint32_t len_x, - const uint32_t* _x, - const IdxT* _in_vals, - uint32_t* _y, - IdxT* _out_vals, - uint32_t* work, - uint32_t* _hints, - bool sort, - uint32_t* _smem) + RAFT_DEVICE_INLINE_FUNCTION void operator()(uint32_t max_topk, + uint32_t topk, + uint32_t len_x, + const uint32_t* _x, + const IdxT* _in_vals, + uint32_t* _y, + IdxT* _out_vals, + uint32_t* work, + uint32_t* _hints, + bool sort, + uint32_t* _smem) { assert(blockDim.x >= V / 4); std::uint8_t* const state = reinterpret_cast(work); - if (max_topk <= 64) { - topk_cta_11_core(topk, - len_x, - _x, - _in_vals, - _y, - _out_vals, - state, - _hints, - sort, - _smem); - } - else if (max_topk <= 128) { - topk_cta_11_core(topk, - len_x, - _x, - _in_vals, - _y, - _out_vals, - state, - _hints, - sort, - _smem); - } - else if (max_topk <= 256) { + if (max_topk <= 256) { topk_cta_11_core(topk, - len_x, - _x, - _in_vals, - _y, - _out_vals, - state, - _hints, - sort, - _smem); - } - else if (max_topk <= 512) { + IdxT>(topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); + } else { + assert(max_topk <= 512); topk_cta_11_core(topk, - len_x, - _x, - _in_vals, - _y, - _out_vals, - state, - _hints, - sort, - _smem); - } - else { - topk_cta_11_core(topk, - len_x, - _x, - _in_vals, - _y, - _out_vals, - state, - _hints, - sort, - _smem); + IdxT>(topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); } } }; From a0be265623c86ff172102fc43c36d40518bab2f3 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Fri, 14 Nov 2025 13:08:29 -0800 Subject: [PATCH 13/25] create non-template wrapper functions to prevent high register pressure of large N path negatively impacting the performance of small N cases --- .../cagra/search_single_cta_kernel-inl.cuh | 217 ++++++++++++++---- .../neighbors/detail/cagra/topk_by_radix.cuh | 85 +++++-- 2 files changed, 235 insertions(+), 67 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 788ae3336e..f9944442fa 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -435,6 +435,127 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_64_false( + float* candidate_distances, // [num_candidates] + std::uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + const std::uint32_t num_itopk, + uint32_t* _smem) +{ + topk_by_bitonic_sort_and_full<64, false, uint32_t>( + candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_128_false( + float* candidate_distances, // [num_candidates] + std::uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + const std::uint32_t num_itopk, + uint32_t* _smem) +{ + topk_by_bitonic_sort_and_full<128, false, uint32_t>( + candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_256_true( + float* candidate_distances, // [num_candidates] + std::uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + const std::uint32_t num_itopk, + uint32_t* _smem) +{ + topk_by_bitonic_sort_and_full<256, true, uint32_t>( + candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_64_false( + float* itopk_distances, // [num_itopk] + uint32_t* itopk_indices, // [num_itopk] + const std::uint32_t num_itopk, + float* candidate_distances, // [num_candidates] + uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + std::uint32_t* work_buf, + const bool first, + uint32_t* _smem) +{ + topk_by_bitonic_sort_and_merge<64, false, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_128_false( + float* itopk_distances, // [num_itopk] + uint32_t* itopk_indices, // [num_itopk] + const std::uint32_t num_itopk, + float* candidate_distances, // [num_candidates] + uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + std::uint32_t* work_buf, + const bool first, + uint32_t* _smem) +{ + topk_by_bitonic_sort_and_merge<128, false, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_false( + float* itopk_distances, // [num_itopk] + uint32_t* itopk_indices, // [num_itopk] + const std::uint32_t num_itopk, + float* candidate_distances, // [num_candidates] + uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + std::uint32_t* work_buf, + const bool first, + uint32_t* _smem) +{ + topk_by_bitonic_sort_and_merge<256, false, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true( + float* itopk_distances, // [num_itopk] + uint32_t* itopk_indices, // [num_itopk] + const std::uint32_t num_itopk, + float* candidate_distances, // [num_candidates] + uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + std::uint32_t* work_buf, + const bool first, + uint32_t* _smem) +{ + topk_by_bitonic_sort_and_merge<512, true, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); +} + template __device__ void topk_by_bitonic_sort_and_merge(float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] @@ -448,73 +569,67 @@ __device__ void topk_by_bitonic_sort_and_merge(float* itopk_distances, // [num_ const bool first, uint32_t* _smem) { + static_assert(std::is_same_v); + // use a non-template wrapper function to avoid co-optimizing for all max_candidatees & max_itopk + // values (which negatively impacts the performance of the functions with smaller max_candidates + // or max_itopk values by over-estimating register pressure) + // The results in candidate_distances/indices are sorted by bitonic sort. assert(blockDim.x >= 64); if (max_candidates <= 64) { - constexpr bool MULTI_WARPS_1 = false; - topk_by_bitonic_sort_and_full<64, MULTI_WARPS_1, IdxT>(candidate_distances, - candidate_indices, - num_candidates, - num_itopk, - static_cast(nullptr)); + topk_by_bitonic_sort_and_full_wrapper_64_false( + candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); } else if (max_candidates <= 128) { - constexpr bool MULTI_WARPS_1 = false; - topk_by_bitonic_sort_and_full<128, MULTI_WARPS_1, IdxT>( + topk_by_bitonic_sort_and_full_wrapper_128_false( candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); } else { - constexpr bool MULTI_WARPS_1 = true; - assert(max_candidates <= 256); - topk_by_bitonic_sort_and_full<256, MULTI_WARPS_1, IdxT>( + topk_by_bitonic_sort_and_full_wrapper_256_true( candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); } // The results sorted above are merged with the internal intermediate top-k // results so far using bitonic merge. if (max_itopk <= 64) { - constexpr bool MULTI_WARPS_2 = false; - topk_by_bitonic_sort_and_merge<64, MULTI_WARPS_2, IdxT>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - static_cast(nullptr)); + topk_by_bitonic_sort_and_merge_wrapper_64_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); } else if (max_itopk <= 128) { - constexpr bool MULTI_WARPS_2 = false; - topk_by_bitonic_sort_and_merge<128, MULTI_WARPS_2, IdxT>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); + topk_by_bitonic_sort_and_merge_wrapper_128_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); } else if (max_itopk <= 256) { - constexpr bool MULTI_WARPS_2 = false; - topk_by_bitonic_sort_and_merge<256, MULTI_WARPS_2, IdxT>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); + topk_by_bitonic_sort_and_merge_wrapper_256_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); } else { assert(max_itopk <= 512); - constexpr bool MULTI_WARPS_2 = true; - topk_by_bitonic_sort_and_merge<512, MULTI_WARPS_2, IdxT>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); + topk_by_bitonic_sort_and_merge_wrapper_512_true(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); } } diff --git a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh index 0d1e679a01..b72165d118 100644 --- a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh @@ -16,9 +16,45 @@ struct topk_by_radix_sort_base { static constexpr uint32_t smem_size(uint32_t max_itopk) { return max_itopk * 2 + 2048 + 8; } }; +RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core_wrapper_256(uint32_t topk, + uint32_t len_x, + const uint32_t* _x, + const uint32_t* _in_vals, + uint32_t* _y, + uint32_t* _out_vals, + uint8_t* state, + uint32_t* _hints, + bool sort, + uint32_t* _smem) +{ + topk_cta_11_core(topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core_wrapper_512(uint32_t topk, + uint32_t len_x, + const uint32_t* _x, + const uint32_t* _in_vals, + uint32_t* _y, + uint32_t* _out_vals, + uint8_t* state, + uint32_t* _hints, + bool sort, + uint32_t* _smem) +{ + topk_cta_11_core(topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); +} + template struct topk_by_radix_sort : topk_by_radix_sort_base { - RAFT_DEVICE_INLINE_FUNCTION void operator()(uint32_t max_topk, + RAFT_DEVICE_INLINE_FUNCTION void operator()(uint32_t max_itopk, uint32_t topk, uint32_t len_x, const uint32_t* _x, @@ -30,21 +66,38 @@ struct topk_by_radix_sort : topk_by_radix_sort_base { bool sort, uint32_t* _smem) { - assert(blockDim.x >= V / 4); - std::uint8_t* const state = reinterpret_cast(work); - if (max_topk <= 256) { - topk_cta_11_core(topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); - } else { - assert(max_topk <= 512); - topk_cta_11_core(topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); + if constexpr (std::is_same_v) { // use a non-template wrapper function to avoid + // co-optimizing the max_itopk <= 256 and + // max_itopk <= 512 cases (which negatively + // impacts the performance of the max_itopk + // <=256 by over-estimating register pressure). + std::uint8_t* const state = reinterpret_cast(work); + if (max_itopk <= 256) { + topk_cta_11_core_wrapper_256( + topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); + } else { + assert(max_itopk <= 512); + topk_cta_11_core_wrapper_512( + topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); + } + } else { // currently, unused + std::uint8_t* const state = reinterpret_cast(work); + if (max_itopk <= 256) { + topk_cta_11_core( + topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); + } else { + assert(max_itopk <= 512); + topk_cta_11_core( + topk, len_x, _x, _in_vals, _y, _out_vals, state, _hints, sort, _smem); + } } } }; From 52141c37cc3fae83cce659b610156d83f2f0b15a Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Fri, 14 Nov 2025 19:36:58 -0800 Subject: [PATCH 14/25] remove unnecessary include statements --- cpp/src/neighbors/detail/cagra/cagra_search.cuh | 1 - cpp/src/neighbors/detail/cagra/search_plan.cuh | 1 - 2 files changed, 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 26e0aafd2d..01577a23d1 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -9,7 +9,6 @@ #include "factory.cuh" #include "sample_filter_utils.cuh" #include "search_plan.cuh" -#include "search_single_cta_inst.cuh" #include #include diff --git a/cpp/src/neighbors/detail/cagra/search_plan.cuh b/cpp/src/neighbors/detail/cagra/search_plan.cuh index ec53ae8cf8..9dcf4aa028 100644 --- a/cpp/src/neighbors/detail/cagra/search_plan.cuh +++ b/cpp/src/neighbors/detail/cagra/search_plan.cuh @@ -10,7 +10,6 @@ #include "compute_distance-ext.cuh" #include #include -// #include "search_single_cta_inst.cuh" // #include "topk_for_cagra/topk.h" #include From 15bdc848b30815606a2636f833681c4174605fed Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Mon, 17 Nov 2025 10:31:58 -0800 Subject: [PATCH 15/25] use smem to reduce register pressure --- .../detail/cagra/search_single_cta.cuh | 10 + .../cagra/search_single_cta_kernel-inl.cuh | 354 ++++++++++++------ .../neighbors/detail/cagra/topk_by_radix.cuh | 9 +- 3 files changed, 244 insertions(+), 129 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh index 986bab0ce0..767ec08dd0 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh @@ -142,6 +142,16 @@ struct search 64); // multi_warps = true if (threadDim.x >= 64) && (itopk_size > 256) max_elements = 512 / 2 /* multi_warps = true */; } + if (num_itopk_candidates > 64) { // use shared memory instead of stack variables (which may + // or may not be placed in registers) for bitonic sort + if (num_itopk_candidates <= 128) { + max_elements = std::max(uint32_t{128}, max_elements); + } else { + static_assert(min_block_size >= + 64); // multi_warps = true if (threadDim.x >= 64) && (itopk_size > 128) + max_elements = std::max(uint32_t{256 / 2 /* multi_warps = true */}, max_elements); + } + } } additional_smem_size += max_block_size * ((max_elements + (raft::warp_size() - 1)) / raft::warp_size()) * diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index f9944442fa..5e279675a3 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -101,7 +101,7 @@ RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parents(std::uint32_t* const termin if (threadIdx.x == 0 && (num_new_parents == 0)) { *terminate_flag = 1; } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] @@ -116,8 +116,17 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( static_assert(MAX_CANDIDATES <= 128); if (warp_id > 0) { return; } constexpr unsigned N = (MAX_CANDIDATES + (raft::warp_size() - 1)) / raft::warp_size(); - float key[N]; - IdxT val[N]; + float* key{}; + IdxT* val{}; + std::conditional_t stack_key{}; + std::conditional_t stack_val{}; + if constexpr (!USE_SMEM) { + key = stack_key; + val = stack_val; + } else { + key = reinterpret_cast(_smem); + val = reinterpret_cast(_smem + N * sizeof(float)); + } /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { unsigned j = lane_id + (raft::warp_size() * i); @@ -144,8 +153,17 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( constexpr unsigned max_candidates_per_warp = (MAX_CANDIDATES + 1) / 2; static_assert(max_candidates_per_warp <= 128); constexpr unsigned N = (max_candidates_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); - float key[N]; - IdxT val[N]; + float* key{}; + IdxT* val{}; + std::conditional_t stack_key{}; + std::conditional_t stack_val{}; + if constexpr (!USE_SMEM) { + key = stack_key; + val = stack_val; + } else { + key = reinterpret_cast(_smem); + val = reinterpret_cast(_smem + N * sizeof(float)); + } if (warp_id < 2) { /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { @@ -208,7 +226,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] @@ -229,9 +247,9 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( constexpr unsigned N = (MAX_ITOPK + (raft::warp_size() - 1)) / raft::warp_size(); float* key{}; IdxT* val{}; - std::conditional_t<(MAX_ITOPK <= 128), float[N], empty_t> stack_key{}; - std::conditional_t<(MAX_ITOPK <= 128), IdxT[N], empty_t> stack_val{}; - if constexpr (MAX_ITOPK <= 128) { + std::conditional_t stack_key{}; + std::conditional_t stack_val{}; + if constexpr (!USE_SMEM) { key = stack_key; val = stack_val; } else { @@ -292,9 +310,9 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( constexpr unsigned N = (max_itopk_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); float* key{}; IdxT* val{}; - std::conditional_t<(max_itopk_per_warp <= 128), float[N], empty_t> stack_key{}; - std::conditional_t<(max_itopk_per_warp <= 128), IdxT[N], empty_t> stack_val{}; - if constexpr (max_itopk_per_warp <= 128) { + std::conditional_t stack_key{}; + std::conditional_t stack_val{}; + if constexpr (!USE_SMEM) { key = stack_key; val = stack_val; } else { @@ -439,32 +457,60 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_64_false( float* candidate_distances, // [num_candidates] std::uint32_t* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, - const std::uint32_t num_itopk, - uint32_t* _smem) + const std::uint32_t num_itopk) { - topk_by_bitonic_sort_and_full<64, false, uint32_t>( - candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); + topk_by_bitonic_sort_and_full<64, false, false, uint32_t>(candidate_distances, + candidate_indices, + num_candidates, + num_itopk, + static_cast(nullptr)); } RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_128_false( + float* candidate_distances, // [num_candidates] + std::uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + const std::uint32_t num_itopk) +{ + topk_by_bitonic_sort_and_full<128, false, false, uint32_t>(candidate_distances, + candidate_indices, + num_candidates, + num_itopk, + static_cast(nullptr)); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_128_false_smem( float* candidate_distances, // [num_candidates] std::uint32_t* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, const std::uint32_t num_itopk, uint32_t* _smem) { - topk_by_bitonic_sort_and_full<128, false, uint32_t>( + topk_by_bitonic_sort_and_full<128, false, true, uint32_t>( candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); } RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_256_true( + float* candidate_distances, // [num_candidates] + std::uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + const std::uint32_t num_itopk) +{ + topk_by_bitonic_sort_and_full<256, true, false, uint32_t>(candidate_distances, + candidate_indices, + num_candidates, + num_itopk, + static_cast(nullptr)); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_256_true_smem( float* candidate_distances, // [num_candidates] std::uint32_t* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, const std::uint32_t num_itopk, uint32_t* _smem) { - topk_by_bitonic_sort_and_full<256, true, uint32_t>( + topk_by_bitonic_sort_and_full<256, true, true, uint32_t>( candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); } @@ -476,21 +522,41 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_64_false uint32_t* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, std::uint32_t* work_buf, - const bool first, - uint32_t* _smem) + const bool first) { - topk_by_bitonic_sort_and_merge<64, false, uint32_t>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); + topk_by_bitonic_sort_and_merge<64, false, false, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + static_cast(nullptr)); } RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_128_false( + float* itopk_distances, // [num_itopk] + uint32_t* itopk_indices, // [num_itopk] + const std::uint32_t num_itopk, + float* candidate_distances, // [num_candidates] + uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t num_candidates, + std::uint32_t* work_buf, + const bool first) +{ + topk_by_bitonic_sort_and_merge<128, false, false, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + static_cast(nullptr)); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_false_smem( float* itopk_distances, // [num_itopk] uint32_t* itopk_indices, // [num_itopk] const std::uint32_t num_itopk, @@ -501,18 +567,18 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_128_fals const bool first, uint32_t* _smem) { - topk_by_bitonic_sort_and_merge<128, false, uint32_t>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); + topk_by_bitonic_sort_and_merge<256, false, true, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); } -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_false( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true_smem( float* itopk_distances, // [num_itopk] uint32_t* itopk_indices, // [num_itopk] const std::uint32_t num_itopk, @@ -523,20 +589,60 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_fals const bool first, uint32_t* _smem) { - topk_by_bitonic_sort_and_merge<256, false, uint32_t>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); + topk_by_bitonic_sort_and_merge<512, true, true, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( + const std::uint32_t max_itopk, + const std::uint32_t num_itopk, + float* candidate_distances, // [num_candidates] + uint32_t* candidate_indices, // [num_candidates] + const std::uint32_t max_candidates, + const std::uint32_t num_candidates, + const bool first, + uint32_t* _smem) +{ + assert(max_candidates <= 256); + + // use a non-template wrapper function to avoid pre-inlining the topk_by_bitonic_sort_and_full + // function (vs post-inlining, this impacts register pressure) + + // The results in candidate_distances/indices are sorted by bitonic sort. + assert(blockDim.x >= 64); + if ((max_itopk > 128) && (max_candidates > 64)) { // store (key, value) pairs in shared memory + if (max_candidates <= 128) { + topk_by_bitonic_sort_and_full_wrapper_128_false_smem( + candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); + } else { + topk_by_bitonic_sort_and_full_wrapper_256_true_smem( + candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); + } + } else { // store (key, value) pairs in stack + if (max_candidates <= 64) { + topk_by_bitonic_sort_and_full_wrapper_64_false( + candidate_distances, candidate_indices, num_candidates, num_itopk); + } else if (max_candidates <= 128) { + topk_by_bitonic_sort_and_full_wrapper_128_false( + candidate_distances, candidate_indices, num_candidates, num_itopk); + } else { + topk_by_bitonic_sort_and_full_wrapper_256_true( + candidate_distances, candidate_indices, num_candidates, num_itopk); + } + } } -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] uint32_t* itopk_indices, // [num_itopk] + const std::uint32_t max_itopk, const std::uint32_t num_itopk, float* candidate_distances, // [num_candidates] uint32_t* candidate_indices, // [num_candidates] @@ -545,92 +651,92 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true const bool first, uint32_t* _smem) { - topk_by_bitonic_sort_and_merge<512, true, uint32_t>(itopk_distances, + // use a non-template wrapper function to avoid pre-inlining the topk_by_bitonic_sort_and_full + // function (vs post-inlining, this impacts register pressure) + + // The results sorted above are merged with the internal intermediate top-k + // results so far using bitonic merge. + if (max_itopk > 128) { // store (key, value) pairs in shared memory + if (max_itopk <= 256) { + topk_by_bitonic_sort_and_merge_wrapper_256_false_smem(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); + } else { + assert(max_itopk <= 512); + topk_by_bitonic_sort_and_merge_wrapper_512_true_smem(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); + } + } else { // store (key, value) pairs in stack + if (max_itopk <= 64) { + topk_by_bitonic_sort_and_merge_wrapper_64_false(itopk_distances, itopk_indices, num_itopk, candidate_distances, candidate_indices, num_candidates, work_buf, - first, - _smem); + first); + } else { + topk_by_bitonic_sort_and_merge_wrapper_128_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); + } + } } template -__device__ void topk_by_bitonic_sort_and_merge(float* itopk_distances, // [num_itopk] - IdxT* itopk_indices, // [num_itopk] - const std::uint32_t max_itopk, - const std::uint32_t num_itopk, - float* candidate_distances, // [num_candidates] - IdxT* candidate_indices, // [num_candidates] - const std::uint32_t max_candidates, - const std::uint32_t num_candidates, - std::uint32_t* work_buf, - const bool first, - uint32_t* _smem) +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( + float* itopk_distances, // [num_itopk] + IdxT* itopk_indices, // [num_itopk] + const std::uint32_t max_itopk, + const std::uint32_t num_itopk, + float* candidate_distances, // [num_candidates] + IdxT* candidate_indices, // [num_candidates] + const std::uint32_t max_candidates, + const std::uint32_t num_candidates, + std::uint32_t* work_buf, + const bool first, + uint32_t* _smem) { static_assert(std::is_same_v); - // use a non-template wrapper function to avoid co-optimizing for all max_candidatees & max_itopk - // values (which negatively impacts the performance of the functions with smaller max_candidates - // or max_itopk values by over-estimating register pressure) - - // The results in candidate_distances/indices are sorted by bitonic sort. - assert(blockDim.x >= 64); - if (max_candidates <= 64) { - topk_by_bitonic_sort_and_full_wrapper_64_false( - candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); - } else if (max_candidates <= 128) { - topk_by_bitonic_sort_and_full_wrapper_128_false( - candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); - } else { - topk_by_bitonic_sort_and_full_wrapper_256_true( - candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); - } - // The results sorted above are merged with the internal intermediate top-k - // results so far using bitonic merge. - if (max_itopk <= 64) { - topk_by_bitonic_sort_and_merge_wrapper_64_false(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); - } else if (max_itopk <= 128) { - topk_by_bitonic_sort_and_merge_wrapper_128_false(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); - } else if (max_itopk <= 256) { - topk_by_bitonic_sort_and_merge_wrapper_256_false(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); - } else { - assert(max_itopk <= 512); - topk_by_bitonic_sort_and_merge_wrapper_512_true(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); - } + topk_by_bitonic_sort_and_full(max_itopk, + num_itopk, + candidate_distances, + candidate_indices, + max_candidates, + num_candidates, + first, + _smem); + + topk_by_bitonic_sort_and_merge(itopk_distances, + itopk_indices, + max_itopk, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first, + _smem); } // This function move the invalid index element to the end of the itopk list. @@ -726,7 +832,7 @@ template -__device__ void search_core( +RAFT_DEVICE_INLINE_FUNCTION void search_core( uintptr_t result_indices_ptr, // [num_queries, top_k] typename DATASET_DESCRIPTOR_T::DISTANCE_T* const result_distances_ptr, // [num_queries, top_k] const std::uint32_t top_k, diff --git a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh index b72165d118..fcd44fcd6d 100644 --- a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh @@ -66,11 +66,10 @@ struct topk_by_radix_sort : topk_by_radix_sort_base { bool sort, uint32_t* _smem) { - if constexpr (std::is_same_v) { // use a non-template wrapper function to avoid - // co-optimizing the max_itopk <= 256 and - // max_itopk <= 512 cases (which negatively - // impacts the performance of the max_itopk - // <=256 by over-estimating register pressure). + if constexpr (std::is_same_v) { // use a non-template wrapper function to avoid + // pre-inlining the topk_cta_11_core function (vs + // post-inlining, this impacts register pressure) std::uint8_t* const state = reinterpret_cast(work); if (max_itopk <= 256) { topk_cta_11_core_wrapper_256( From 5a117a58aae4a6cf9443a633d31d9f4954bfa10f Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Mon, 17 Nov 2025 15:40:00 -0800 Subject: [PATCH 16/25] fix build error after pulling new updates --- .../detail/cagra/search_single_cta_kernel-inl.cuh | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 8291b7874e..2085591d78 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -1518,9 +1518,6 @@ struct search_kernel_config { // Radix-based topk is used return dispatch_kernel; } - THROW("No kernel for parameters itopk_size %u, num_itopk_candidates %u", - itopk_size, - num_itopk_candidates); } }; @@ -2375,6 +2372,7 @@ void select_and_run( } uint32_t max_itopk{}; + assert(ps.itopk_size <= 512); if (num_itopk_candidates <= 256) { // bitonic sort if (ps.itopk_size <= 64) { max_itopk = 64; @@ -2382,18 +2380,14 @@ void select_and_run( max_itopk = 128; } else if (ps.itopk_size <= 256) { max_itopk = 256; - } else if (ps.itopk_size <= 512) { - max_itopk = 512; } else { - THROW("No kernel for parameter itopk_size %u", ps.itopk_size); + max_itopk = 512; } } else { // radix sort if (ps.itopk_size <= 256) { max_itopk = 256; - } else if (ps.itopk_size <= 512) { - max_itopk = 512; } else { - THROW("No kernel for parameter itopk_size %u", ps.itopk_size); + max_itopk = 512; } } From 4cd6f8f0937553ff97519803f0cadf9677adcacb Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Mon, 17 Nov 2025 20:25:55 -0800 Subject: [PATCH 17/25] undo using shared memory to store key, value pairs --- .../detail/cagra/search_single_cta.cuh | 29 +- .../cagra/search_single_cta_kernel-inl.cuh | 287 ++++++------------ 2 files changed, 87 insertions(+), 229 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh index 767ec08dd0..b656236164 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh @@ -129,34 +129,7 @@ struct search sizeof(std::uint32_t) * topk_ws_size + sizeof(std::uint32_t); std::uint32_t additional_smem_size = 0; - if (num_itopk_candidates <= 256) { // bitonic sort - // Tentatively calculate the required shared memory size when bitonic sort based topk is used, - // assuming the block size is the maximum. - uint32_t max_elements{0}; - if (itopk_size > 128) { // use shared memory instead of stack variables (which may or may not - // be placed in registers) for bitonic sort - if (itopk_size <= 256) { - max_elements = 256; - } else { - static_assert(min_block_size >= - 64); // multi_warps = true if (threadDim.x >= 64) && (itopk_size > 256) - max_elements = 512 / 2 /* multi_warps = true */; - } - if (num_itopk_candidates > 64) { // use shared memory instead of stack variables (which may - // or may not be placed in registers) for bitonic sort - if (num_itopk_candidates <= 128) { - max_elements = std::max(uint32_t{128}, max_elements); - } else { - static_assert(min_block_size >= - 64); // multi_warps = true if (threadDim.x >= 64) && (itopk_size > 128) - max_elements = std::max(uint32_t{256 / 2 /* multi_warps = true */}, max_elements); - } - } - } - additional_smem_size += max_block_size * - ((max_elements + (raft::warp_size() - 1)) / raft::warp_size()) * - (sizeof(float) + sizeof(INDEX_T)); - } else { // radix sort + if (num_itopk_candidates > 256) { // radix sort // Tentatively calculate the required shared memory size when radix sort based topk is used, // assuming the block size is the maximum. if (itopk_size <= 256) { diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 2085591d78..fcb6f4d233 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -101,13 +101,12 @@ RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parents(std::uint32_t* const termin if (threadIdx.x == 0 && (num_new_parents == 0)) { *terminate_flag = 1; } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( float* candidate_distances, // [num_candidates] IdxT* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, - const std::uint32_t num_itopk, - uint32_t* _smem) + const std::uint32_t num_itopk) { const unsigned lane_id = threadIdx.x % raft::warp_size(); const unsigned warp_id = threadIdx.x / raft::warp_size(); @@ -116,17 +115,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( static_assert(MAX_CANDIDATES <= 128); if (warp_id > 0) { return; } constexpr unsigned N = (MAX_CANDIDATES + (raft::warp_size() - 1)) / raft::warp_size(); - float* key{}; - IdxT* val{}; - std::conditional_t stack_key{}; - std::conditional_t stack_val{}; - if constexpr (!USE_SMEM) { - key = stack_key; - val = stack_val; - } else { - key = reinterpret_cast(_smem); - val = reinterpret_cast(_smem + N * sizeof(float)); - } + float key[N]; + IdxT val[N]; /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { unsigned j = lane_id + (raft::warp_size() * i); @@ -153,17 +143,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( constexpr unsigned max_candidates_per_warp = (MAX_CANDIDATES + 1) / 2; static_assert(max_candidates_per_warp <= 128); constexpr unsigned N = (max_candidates_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); - float* key{}; - IdxT* val{}; - std::conditional_t stack_key{}; - std::conditional_t stack_val{}; - if constexpr (!USE_SMEM) { - key = stack_key; - val = stack_val; - } else { - key = reinterpret_cast(_smem); - val = reinterpret_cast(_smem + N * sizeof(float)); - } + float key[N]; + IdxT val[N]; if (warp_id < 2) { /* Candidates -> Reg */ for (unsigned i = 0; i < N; i++) { @@ -226,7 +207,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } -template +template RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( float* itopk_distances, // [num_itopk] IdxT* itopk_indices, // [num_itopk] @@ -235,8 +216,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( IdxT* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, std::uint32_t* work_buf, - const bool first, - uint32_t* _smem) + const bool first) { const unsigned lane_id = threadIdx.x % raft::warp_size(); const unsigned warp_id = threadIdx.x / raft::warp_size(); @@ -245,17 +225,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( if constexpr (!MULTI_WARPS) { if (warp_id > 0) { return; } constexpr unsigned N = (MAX_ITOPK + (raft::warp_size() - 1)) / raft::warp_size(); - float* key{}; - IdxT* val{}; - std::conditional_t stack_key{}; - std::conditional_t stack_val{}; - if constexpr (!USE_SMEM) { - key = stack_key; - val = stack_val; - } else { - key = reinterpret_cast(_smem); - val = reinterpret_cast(_smem + N * sizeof(float)); - } + float key[N]; + IdxT val[N]; if (first) { /* Load itopk results */ for (unsigned i = 0; i < N; i++) { @@ -308,17 +279,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( // Use two warps (64 threads) or more constexpr unsigned max_itopk_per_warp = (MAX_ITOPK + 1) / 2; constexpr unsigned N = (max_itopk_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); - float* key{}; - IdxT* val{}; - std::conditional_t stack_key{}; - std::conditional_t stack_val{}; - if constexpr (!USE_SMEM) { - key = stack_key; - val = stack_val; - } else { - key = reinterpret_cast(_smem); - val = reinterpret_cast(_smem + N * sizeof(float)); - } + float key[N]; + IdxT val[N]; if (first) { /* Load itop results (not sorted) */ if (warp_id < 2) { @@ -459,11 +421,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_64_false( const std::uint32_t num_candidates, const std::uint32_t num_itopk) { - topk_by_bitonic_sort_and_full<64, false, false, uint32_t>(candidate_distances, - candidate_indices, - num_candidates, - num_itopk, - static_cast(nullptr)); + topk_by_bitonic_sort_and_full<64, false, uint32_t>( + candidate_distances, candidate_indices, num_candidates, num_itopk); } RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_128_false( @@ -472,22 +431,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_128_false const std::uint32_t num_candidates, const std::uint32_t num_itopk) { - topk_by_bitonic_sort_and_full<128, false, false, uint32_t>(candidate_distances, - candidate_indices, - num_candidates, - num_itopk, - static_cast(nullptr)); -} - -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_128_false_smem( - float* candidate_distances, // [num_candidates] - std::uint32_t* candidate_indices, // [num_candidates] - const std::uint32_t num_candidates, - const std::uint32_t num_itopk, - uint32_t* _smem) -{ - topk_by_bitonic_sort_and_full<128, false, true, uint32_t>( - candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); + topk_by_bitonic_sort_and_full<128, false, uint32_t>( + candidate_distances, candidate_indices, num_candidates, num_itopk); } RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_256_true( @@ -496,22 +441,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_256_true( const std::uint32_t num_candidates, const std::uint32_t num_itopk) { - topk_by_bitonic_sort_and_full<256, true, false, uint32_t>(candidate_distances, - candidate_indices, - num_candidates, - num_itopk, - static_cast(nullptr)); -} - -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_256_true_smem( - float* candidate_distances, // [num_candidates] - std::uint32_t* candidate_indices, // [num_candidates] - const std::uint32_t num_candidates, - const std::uint32_t num_itopk, - uint32_t* _smem) -{ - topk_by_bitonic_sort_and_full<256, true, true, uint32_t>( - candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); + topk_by_bitonic_sort_and_full<256, true, uint32_t>( + candidate_distances, candidate_indices, num_candidates, num_itopk); } RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_64_false( @@ -524,15 +455,14 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_64_false std::uint32_t* work_buf, const bool first) { - topk_by_bitonic_sort_and_merge<64, false, false, uint32_t>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - static_cast(nullptr)); + topk_by_bitonic_sort_and_merge<64, false, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); } RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_128_false( @@ -545,18 +475,17 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_128_fals std::uint32_t* work_buf, const bool first) { - topk_by_bitonic_sort_and_merge<128, false, false, uint32_t>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - static_cast(nullptr)); + topk_by_bitonic_sort_and_merge<128, false, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); } -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_false_smem( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_false( float* itopk_distances, // [num_itopk] uint32_t* itopk_indices, // [num_itopk] const std::uint32_t num_itopk, @@ -567,18 +496,17 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_fals const bool first, uint32_t* _smem) { - topk_by_bitonic_sort_and_merge<256, false, true, uint32_t>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); + topk_by_bitonic_sort_and_merge<256, false, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); } -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true_smem( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true( float* itopk_distances, // [num_itopk] uint32_t* itopk_indices, // [num_itopk] const std::uint32_t num_itopk, @@ -589,15 +517,14 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true const bool first, uint32_t* _smem) { - topk_by_bitonic_sort_and_merge<512, true, true, uint32_t>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); + topk_by_bitonic_sort_and_merge<512, true, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); } RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( @@ -607,8 +534,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( uint32_t* candidate_indices, // [num_candidates] const std::uint32_t max_candidates, const std::uint32_t num_candidates, - const bool first, - uint32_t* _smem) + const bool first) { assert(max_candidates <= 256); @@ -617,25 +543,15 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( // The results in candidate_distances/indices are sorted by bitonic sort. assert(blockDim.x >= 64); - if ((max_itopk > 128) && (max_candidates > 64)) { // store (key, value) pairs in shared memory - if (max_candidates <= 128) { - topk_by_bitonic_sort_and_full_wrapper_128_false_smem( - candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); - } else { - topk_by_bitonic_sort_and_full_wrapper_256_true_smem( - candidate_distances, candidate_indices, num_candidates, num_itopk, _smem); - } - } else { // store (key, value) pairs in stack - if (max_candidates <= 64) { - topk_by_bitonic_sort_and_full_wrapper_64_false( - candidate_distances, candidate_indices, num_candidates, num_itopk); - } else if (max_candidates <= 128) { - topk_by_bitonic_sort_and_full_wrapper_128_false( - candidate_distances, candidate_indices, num_candidates, num_itopk); - } else { - topk_by_bitonic_sort_and_full_wrapper_256_true( - candidate_distances, candidate_indices, num_candidates, num_itopk); - } + if (max_candidates <= 64) { + topk_by_bitonic_sort_and_full_wrapper_64_false( + candidate_distances, candidate_indices, num_candidates, num_itopk); + } else if (max_candidates <= 128) { + topk_by_bitonic_sort_and_full_wrapper_128_false( + candidate_distances, candidate_indices, num_candidates, num_itopk); + } else { + topk_by_bitonic_sort_and_full_wrapper_256_true( + candidate_distances, candidate_indices, num_candidates, num_itopk); } } @@ -648,57 +564,31 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( uint32_t* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, std::uint32_t* work_buf, - const bool first, - uint32_t* _smem) + const bool first) { // use a non-template wrapper function to avoid pre-inlining the topk_by_bitonic_sort_and_full // function (vs post-inlining, this impacts register pressure) // The results sorted above are merged with the internal intermediate top-k // results so far using bitonic merge. - if (max_itopk > 128) { // store (key, value) pairs in shared memory - if (max_itopk <= 256) { - topk_by_bitonic_sort_and_merge_wrapper_256_false_smem(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); - } else { - assert(max_itopk <= 512); - topk_by_bitonic_sort_and_merge_wrapper_512_true_smem(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first, - _smem); - } - } else { // store (key, value) pairs in stack - if (max_itopk <= 64) { - topk_by_bitonic_sort_and_merge_wrapper_64_false(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first); - } else { - topk_by_bitonic_sort_and_merge_wrapper_128_false(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first); - } + if (max_itopk <= 64) { + topk_by_bitonic_sort_and_merge_wrapper_64_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); + } else { + topk_by_bitonic_sort_and_merge_wrapper_128_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); } } @@ -713,8 +603,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( const std::uint32_t max_candidates, const std::uint32_t num_candidates, std::uint32_t* work_buf, - const bool first, - uint32_t* _smem) + const bool first) { static_assert(std::is_same_v); @@ -724,8 +613,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( candidate_indices, max_candidates, num_candidates, - first, - _smem); + first); topk_by_bitonic_sort_and_merge(itopk_distances, itopk_indices, @@ -735,8 +623,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( candidate_indices, num_candidates, work_buf, - first, - _smem); + first); } // This function move the invalid index element to the end of the itopk list. @@ -1014,8 +901,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( max_candidates, search_width * graph_degree, topk_ws, - (iter == 0), - smem_work_ptr); + (iter == 0)); __syncthreads(); _CLK_REC(clk_topk); } else { @@ -1180,8 +1066,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( max_candidates, search_width * graph_degree, topk_ws, - (iter == 0), - smem_work_ptr); + (iter == 0)); } __syncthreads(); } From c8a9673b0712d89ba1c08cc476fa26a6be69a1f9 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Mon, 17 Nov 2025 21:26:28 -0800 Subject: [PATCH 18/25] delete dead code --- cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 2 -- 1 file changed, 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index fcb6f4d233..ae78647314 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -55,8 +55,6 @@ namespace cuvs::neighbors::cagra::detail { namespace single_cta_search { -struct empty_t {}; - // #define _CLK_BREAKDOWN template From 77d657c10a9cac2c82547dc259f13d8809638027 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Mon, 17 Nov 2025 23:59:57 -0800 Subject: [PATCH 19/25] fix an error --- .../cagra/search_single_cta_kernel-inl.cuh | 27 +++++++++++++++---- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index ae78647314..d9456263e0 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -491,8 +491,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_fals uint32_t* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, std::uint32_t* work_buf, - const bool first, - uint32_t* _smem) + const bool first) { topk_by_bitonic_sort_and_merge<256, false, uint32_t>(itopk_distances, itopk_indices, @@ -512,8 +511,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true uint32_t* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, std::uint32_t* work_buf, - const bool first, - uint32_t* _smem) + const bool first) { topk_by_bitonic_sort_and_merge<512, true, uint32_t>(itopk_distances, itopk_indices, @@ -578,7 +576,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( num_candidates, work_buf, first); - } else { + } else if (max_itopk <= 128) { topk_by_bitonic_sort_and_merge_wrapper_128_false(itopk_distances, itopk_indices, num_itopk, @@ -587,6 +585,25 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( num_candidates, work_buf, first); + } else if (max_itopk <= 256) { + topk_by_bitonic_sort_and_merge_wrapper_256_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); + } else { + assert(max_itopk <= 512); + topk_by_bitonic_sort_and_merge_wrapper_512_true(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); } } From efcf48980778b394b4a5c65622ceec55585a0414 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Tue, 18 Nov 2025 02:58:46 -0800 Subject: [PATCH 20/25] tweak register pressure --- .../cagra/search_single_cta_kernel-inl.cuh | 56 ++++++++++--------- 1 file changed, 29 insertions(+), 27 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index d9456263e0..9285b799f0 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -567,33 +567,35 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( // The results sorted above are merged with the internal intermediate top-k // results so far using bitonic merge. - if (max_itopk <= 64) { - topk_by_bitonic_sort_and_merge_wrapper_64_false(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first); - } else if (max_itopk <= 128) { - topk_by_bitonic_sort_and_merge_wrapper_128_false(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first); - } else if (max_itopk <= 256) { - topk_by_bitonic_sort_and_merge_wrapper_256_false(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first); + if (max_itopk <= 256) { + if (max_itopk <= 64) { + topk_by_bitonic_sort_and_merge_wrapper_64_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); + } else if (max_itopk <= 128) { + topk_by_bitonic_sort_and_merge_wrapper_128_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); + } else { + topk_by_bitonic_sort_and_merge_wrapper_256_false(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); + } } else { assert(max_itopk <= 512); topk_by_bitonic_sort_and_merge_wrapper_512_true(itopk_distances, From bde5ac515ca3ce12981270ff1643eaba8b7383f8 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Tue, 18 Nov 2025 10:49:49 -0800 Subject: [PATCH 21/25] copyright year --- cpp/src/neighbors/detail/cagra/bitonic.hpp | 2 +- cpp/src/neighbors/detail/cagra/topk_by_radix.cuh | 2 +- cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/bitonic.hpp b/cpp/src/neighbors/detail/cagra/bitonic.hpp index 58cd83a5b3..ace2d8ae20 100644 --- a/cpp/src/neighbors/detail/cagra/bitonic.hpp +++ b/cpp/src/neighbors/detail/cagra/bitonic.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh index fcd44fcd6d..7b0e3dbcad 100644 --- a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh index 1e9bd86389..1565faba2d 100644 --- a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once From 901bad6926126a6a6e9836be6248163a252a4bcb Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Sat, 22 Nov 2025 19:43:44 -0800 Subject: [PATCH 22/25] undo conditional unrolling in bitonic.hpp (this significantly slows down (16-17%) max_itopk=512 with batch size 10000 cases --- cpp/src/neighbors/detail/cagra/bitonic.hpp | 11 +- .../detail/cagra/search_single_cta.cuh | 2 +- .../cagra/search_single_cta_kernel-inl.cuh | 334 ++++++++---------- 3 files changed, 160 insertions(+), 187 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/bitonic.hpp b/cpp/src/neighbors/detail/cagra/bitonic.hpp index ace2d8ae20..7fd81f0a11 100644 --- a/cpp/src/neighbors/detail/cagra/bitonic.hpp +++ b/cpp/src/neighbors/detail/cagra/bitonic.hpp @@ -226,16 +226,9 @@ RAFT_DEVICE_INLINE_FUNCTION void warp_merge(K k[N], V v[N], unsigned range, cons template RAFT_DEVICE_INLINE_FUNCTION void warp_sort(K k[N], V v[N], const bool asc = true) { - if constexpr (N < 8) { #pragma unroll - for (std::uint32_t range = 1; range <= warp_size; range <<= 1) { - warp_merge(k, v, range, asc); - } - } else { -#pragma unroll 1 - for (std::uint32_t range = 1; range <= warp_size; range <<= 1) { - warp_merge(k, v, range, asc); - } + for (std::uint32_t range = 1; range <= warp_size; range <<= 1) { + warp_merge(k, v, range, asc); } } diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh index b656236164..b2b4670964 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh @@ -117,7 +117,7 @@ struct search // // Determine the thread block size // - constexpr unsigned min_block_size = 64; // 32 or 64 + constexpr unsigned min_block_size = 64; constexpr unsigned min_block_size_radix = 256; constexpr unsigned max_block_size = 1024; // diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 9285b799f0..76a45dff81 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -57,7 +57,7 @@ namespace single_cta_search { // #define _CLK_BREAKDOWN -template +template RAFT_DEVICE_INLINE_FUNCTION void pickup_next_parents(std::uint32_t* const terminate_flag, INDEX_T* const next_parent_indices, INDEX_T* const internal_topk_indices, @@ -110,7 +110,6 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( const unsigned warp_id = threadIdx.x / raft::warp_size(); static_assert(MAX_CANDIDATES <= 256); if constexpr (!MULTI_WARPS) { - static_assert(MAX_CANDIDATES <= 128); if (warp_id > 0) { return; } constexpr unsigned N = (MAX_CANDIDATES + (raft::warp_size() - 1)) / raft::warp_size(); float key[N]; @@ -137,6 +136,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( } } } else { + assert(blockDim.x >= 64); // Use two warps (64 threads) constexpr unsigned max_candidates_per_warp = (MAX_CANDIDATES + 1) / 2; static_assert(max_candidates_per_warp <= 128); @@ -221,6 +221,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( static_assert(MAX_ITOPK <= 512); if constexpr (!MULTI_WARPS) { + static_assert(MAX_ITOPK <= 256); if (warp_id > 0) { return; } constexpr unsigned N = (MAX_ITOPK + (raft::warp_size() - 1)) / raft::warp_size(); float key[N]; @@ -274,6 +275,8 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( } } } else { + static_assert(MAX_ITOPK == 512); + assert(blockDim.x >= 64); // Use two warps (64 threads) or more constexpr unsigned max_itopk_per_warp = (MAX_ITOPK + 1) / 2; constexpr unsigned N = (max_itopk_per_warp + (raft::warp_size() - 1)) / raft::warp_size(); @@ -433,13 +436,13 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_128_false candidate_distances, candidate_indices, num_candidates, num_itopk); } -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_256_true( +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full_wrapper_256_false( float* candidate_distances, // [num_candidates] std::uint32_t* candidate_indices, // [num_candidates] const std::uint32_t num_candidates, const std::uint32_t num_itopk) { - topk_by_bitonic_sort_and_full<256, true, uint32_t>( + topk_by_bitonic_sort_and_full<256, false, uint32_t>( candidate_distances, candidate_indices, num_candidates, num_itopk); } @@ -503,42 +506,26 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_256_fals first); } -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge_wrapper_512_true( - float* itopk_distances, // [num_itopk] - uint32_t* itopk_indices, // [num_itopk] - const std::uint32_t num_itopk, - float* candidate_distances, // [num_candidates] - uint32_t* candidate_indices, // [num_candidates] - const std::uint32_t num_candidates, - std::uint32_t* work_buf, - const bool first) -{ - topk_by_bitonic_sort_and_merge<512, true, uint32_t>(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first); -} - -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( +template +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( + float* itopk_distances, // [num_itopk] + IdxT* itopk_indices, // [num_itopk] const std::uint32_t max_itopk, const std::uint32_t num_itopk, - float* candidate_distances, // [num_candidates] - uint32_t* candidate_indices, // [num_candidates] + float* candidate_distances, // [num_candidates] + IdxT* candidate_indices, // [num_candidates] const std::uint32_t max_candidates, const std::uint32_t num_candidates, + std::uint32_t* work_buf, const bool first) { + static_assert(std::is_same_v); + assert(max_itopk <= 512); assert(max_candidates <= 256); + assert(!MULTI_WARPS || blockDim.x >= 64); // use a non-template wrapper function to avoid pre-inlining the topk_by_bitonic_sort_and_full // function (vs post-inlining, this impacts register pressure) - - // The results in candidate_distances/indices are sorted by bitonic sort. - assert(blockDim.x >= 64); if (max_candidates <= 64) { topk_by_bitonic_sort_and_full_wrapper_64_false( candidate_distances, candidate_indices, num_candidates, num_itopk); @@ -546,28 +533,14 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_full( topk_by_bitonic_sort_and_full_wrapper_128_false( candidate_distances, candidate_indices, num_candidates, num_itopk); } else { - topk_by_bitonic_sort_and_full_wrapper_256_true( + topk_by_bitonic_sort_and_full_wrapper_256_false( candidate_distances, candidate_indices, num_candidates, num_itopk); } -} - -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( - float* itopk_distances, // [num_itopk] - uint32_t* itopk_indices, // [num_itopk] - const std::uint32_t max_itopk, - const std::uint32_t num_itopk, - float* candidate_distances, // [num_candidates] - uint32_t* candidate_indices, // [num_candidates] - const std::uint32_t num_candidates, - std::uint32_t* work_buf, - const bool first) -{ - // use a non-template wrapper function to avoid pre-inlining the topk_by_bitonic_sort_and_full - // function (vs post-inlining, this impacts register pressure) - // The results sorted above are merged with the internal intermediate top-k - // results so far using bitonic merge. - if (max_itopk <= 256) { + if constexpr (!MULTI_WARPS) { + assert(max_itopk <= 256); + // use a non-template wrapper function to avoid pre-inlining the topk_by_bitonic_sort_and_merge + // function (vs post-inlining, this impacts register pressure) if (max_itopk <= 64) { topk_by_bitonic_sort_and_merge_wrapper_64_false(itopk_distances, itopk_indices, @@ -597,52 +570,18 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( first); } } else { - assert(max_itopk <= 512); - topk_by_bitonic_sort_and_merge_wrapper_512_true(itopk_distances, - itopk_indices, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first); + assert(max_itopk > 256); + topk_by_bitonic_sort_and_merge<512, MULTI_WARPS, uint32_t>(itopk_distances, + itopk_indices, + num_itopk, + candidate_distances, + candidate_indices, + num_candidates, + work_buf, + first); } } -template -RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_and_merge( - float* itopk_distances, // [num_itopk] - IdxT* itopk_indices, // [num_itopk] - const std::uint32_t max_itopk, - const std::uint32_t num_itopk, - float* candidate_distances, // [num_candidates] - IdxT* candidate_indices, // [num_candidates] - const std::uint32_t max_candidates, - const std::uint32_t num_candidates, - std::uint32_t* work_buf, - const bool first) -{ - static_assert(std::is_same_v); - - topk_by_bitonic_sort_and_full(max_itopk, - num_itopk, - candidate_distances, - candidate_indices, - max_candidates, - num_candidates, - first); - - topk_by_bitonic_sort_and_merge(itopk_distances, - itopk_indices, - max_itopk, - num_itopk, - candidate_distances, - candidate_indices, - num_candidates, - work_buf, - first); -} - // This function move the invalid index element to the end of the itopk list. // Require : array_length % 32 == 0 && The invalid entry is only one. template @@ -732,7 +671,8 @@ RAFT_DEVICE_INLINE_FUNCTION void hashmap_restore(INDEX_T* const hashmap_ptr, * @param small_hash_reset_interval Interval for resetting the small hash. * @param query_id sequential id of the query in the batch */ -template @@ -867,8 +807,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( // if max_candidates is greater than 128, the first operation uses two warps; // if max_itopk is greater than 256, the second operation used two warps. assert(blockDim.x >= 64); - const bool multi_warps_1 = (max_candidates > 128) ? true : false; - const bool multi_warps_2 = (max_itopk > 256) ? true : false; + const bool bitonic_sort_and_full_multi_warps = (max_candidates > 128) ? true : false; // reset small-hash table. if ((iter + 1) % small_hash_reset_interval == 0) { @@ -881,13 +820,13 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( if (blockDim.x == 32) { hash_start_tid = 0; } else if (blockDim.x == 64) { - if (multi_warps_1 || multi_warps_2) { + if (bitonic_sort_and_full_multi_warps || BITONIC_SORT_AND_MERGE_MULTI_WARPS) { hash_start_tid = 0; } else { hash_start_tid = 32; } } else { - if (multi_warps_1 || multi_warps_2) { + if (bitonic_sort_and_full_multi_warps || BITONIC_SORT_AND_MERGE_MULTI_WARPS) { hash_start_tid = 64; } else { hash_start_tid = 32; @@ -909,16 +848,17 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( if (threadIdx.x == 0) { *terminate_flag = 0; } } - topk_by_bitonic_sort_and_merge(result_distances_buffer, - result_indices_buffer, - max_itopk, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - max_candidates, - search_width * graph_degree, - topk_ws, - (iter == 0)); + topk_by_bitonic_sort_and_merge( + result_distances_buffer, + result_indices_buffer, + max_itopk, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + max_candidates, + search_width * graph_degree, + topk_ws, + (iter == 0)); __syncthreads(); _CLK_REC(clk_topk); } else { @@ -1074,16 +1014,17 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( // candidate list. if (top_k > internal_topk || result_indices_buffer[top_k - 1] == invalid_index) { __syncthreads(); - topk_by_bitonic_sort_and_merge(result_distances_buffer, - result_indices_buffer, - max_itopk, - internal_topk, - result_distances_buffer + internal_topk, - result_indices_buffer + internal_topk, - max_candidates, - search_width * graph_degree, - topk_ws, - (iter == 0)); + topk_by_bitonic_sort_and_merge( + result_distances_buffer, + result_indices_buffer, + max_itopk, + internal_topk, + result_distances_buffer + internal_topk, + result_indices_buffer + internal_topk, + max_candidates, + search_width * graph_degree, + topk_ws, + (iter == 0)); } __syncthreads(); } @@ -1153,7 +1094,8 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( #endif } -template @@ -1185,32 +1127,35 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( SAMPLE_FILTER_T sample_filter) { const auto query_id = blockIdx.y; - search_core( - result_indices_ptr, - result_distances_ptr, - top_k, - dataset_desc, - queries_ptr, - knn_graph, - graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hashmap_ptr, - max_candidates, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - sample_filter); + search_core(result_indices_ptr, + result_distances_ptr, + top_k, + dataset_desc, + queries_ptr, + knn_graph, + graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + sample_filter); } // To make sure we avoid false sharing on both CPU and GPU, we enforce cache line size to the @@ -1272,7 +1217,8 @@ constexpr auto is_worker_busy(worker_handle_t::handle_t h) -> bool return (h != kWaitForWork) && (h != kNoMoreWork); } -template @@ -1343,32 +1289,35 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel_p( auto query_id = worker_data.value.query_id; // work phase - search_core( - result_indices_ptr, - result_distances_ptr, - top_k, - dataset_desc, - queries_ptr, - knn_graph, - graph_degree, - source_indices_ptr, - num_distilation, - rand_xor_mask, - seed_ptr, - num_seeds, - visited_hashmap_ptr, - max_candidates, - max_itopk, - internal_topk, - search_width, - min_iteration, - max_iteration, - num_executed_iterations, - hash_bitlen, - small_hash_bitlen, - small_hash_reset_interval, - query_id, - sample_filter); + search_core(result_indices_ptr, + result_distances_ptr, + top_k, + dataset_desc, + queries_ptr, + knn_graph, + graph_degree, + source_indices_ptr, + num_distilation, + rand_xor_mask, + seed_ptr, + num_seeds, + visited_hashmap_ptr, + max_candidates, + max_itopk, + internal_topk, + search_width, + min_iteration, + max_iteration, + num_executed_iterations, + hash_bitlen, + small_hash_bitlen, + small_hash_reset_interval, + query_id, + sample_filter); // make sure all writes are visible even for the host // (e.g. when result buffers are in pinned memory) @@ -1386,18 +1335,25 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel_p( } template auto dispatch_kernel = []() { + static_assert(TOPK_BY_BITONIC_SORT || !BITONIC_SORT_AND_MERGE_MULTI_WARPS); if constexpr (Persistent) { return search_kernel_p; } else { - return search_kernel; + return search_kernel; } }(); @@ -1406,8 +1362,12 @@ template struct search_kernel_config { - using kernel_t = - decltype(dispatch_kernel); + using kernel_t = decltype(dispatch_kernel); static auto choose_itopk_and_mx_candidates(unsigned itopk_size, unsigned num_itopk_candidates, @@ -1415,10 +1375,30 @@ struct search_kernel_config { { assert(itopk_size <= 512); if (num_itopk_candidates <= 256) { - return dispatch_kernel; + if (itopk_size <= 256) { + return dispatch_kernel; + } else { + assert(block_size >= 64); + return dispatch_kernel; + } } else { // Radix-based topk is used - return dispatch_kernel; + return dispatch_kernel; } } }; From 4f7ef48a84942cfaef1847faed211a91373efafc Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Mon, 24 Nov 2025 09:31:23 -0800 Subject: [PATCH 23/25] final performance tweak --- .../cagra/search_multi_cta_kernel-inl.cuh | 58 ++++++++++++++++--- .../detail/cagra/topk_for_cagra/topk_core.cuh | 1 + 2 files changed, 50 insertions(+), 9 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 4d5fffb972..44aa4c9601 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -142,6 +142,30 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort(float* distances, // [num } } +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_wrapper_64( + float* distances, // [num_elements] + uint32_t* indices, // [num_elements] + const uint32_t num_elements) +{ + topk_by_bitonic_sort<64, uint32_t>(distances, indices, num_elements); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_wrapper_128( + float* distances, // [num_elements] + uint32_t* indices, // [num_elements] + const uint32_t num_elements) +{ + topk_by_bitonic_sort<128, uint32_t>(distances, indices, num_elements); +} + +RAFT_DEVICE_INLINE_FUNCTION void topk_by_bitonic_sort_wrapper_256( + float* distances, // [num_elements] + uint32_t* indices, // [num_elements] + const uint32_t num_elements) +{ + topk_by_bitonic_sort<256, uint32_t>(distances, indices, num_elements); +} + // // multiple CTAs per single query // @@ -266,16 +290,32 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( _CLK_START(); if (threadIdx.x < 32) { // [1st warp] Topk with bitonic sort - if (max_elements <= 64) { - topk_by_bitonic_sort<64, INDEX_T>( - result_distances_buffer, result_indices_buffer, result_buffer_size_32); - } else if (max_elements <= 128) { - topk_by_bitonic_sort<128, INDEX_T>( - result_distances_buffer, result_indices_buffer, result_buffer_size_32); + if constexpr (std::is_same_v) { + // use a non-template wrapper function to avoid pre-inlining the topk_by_bitonic_sort + // function (vs post-inlining, this impacts register pressure) + if (max_elements <= 64) { + topk_by_bitonic_sort_wrapper_64( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } else if (max_elements <= 128) { + topk_by_bitonic_sort_wrapper_128( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } else { + assert(max_elements <= 256); + topk_by_bitonic_sort_wrapper_256( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } } else { - assert(max_elements <= 256); - topk_by_bitonic_sort<256, INDEX_T>( - result_distances_buffer, result_indices_buffer, result_buffer_size_32); + if (max_elements <= 64) { + topk_by_bitonic_sort<64, INDEX_T>( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } else if (max_elements <= 128) { + topk_by_bitonic_sort<128, INDEX_T>( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } else { + assert(max_elements <= 256); + topk_by_bitonic_sort<256, INDEX_T>( + result_distances_buffer, result_indices_buffer, result_buffer_size_32); + } } } __syncthreads(); diff --git a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh index 1565faba2d..97ea7dc236 100644 --- a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh @@ -643,6 +643,7 @@ RAFT_DEVICE_INLINE_FUNCTION void topk_cta_11_core(uint32_t topk, // // Search for the maximum threshold that satisfies "(x < threshold).sum() <= topk". // +#pragma unroll for (int j = 0; j < 3; j += 1) { uint32_t num_bins; uint32_t shift; From 64595f92b433de56208ab7eef4f481504948e2c0 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Fri, 5 Dec 2025 10:43:40 -0800 Subject: [PATCH 24/25] copyright year --- cpp/src/neighbors/detail/cagra/bitonic.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/src/neighbors/detail/cagra/bitonic.hpp b/cpp/src/neighbors/detail/cagra/bitonic.hpp index 7fd81f0a11..66c726c71a 100644 --- a/cpp/src/neighbors/detail/cagra/bitonic.hpp +++ b/cpp/src/neighbors/detail/cagra/bitonic.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once From 3dc5569a1373c6c39bd7dd23ebb99d7cb593bd81 Mon Sep 17 00:00:00 2001 From: Seunghwa Kang Date: Wed, 7 Jan 2026 07:34:10 -0800 Subject: [PATCH 25/25] copyright year --- cpp/src/neighbors/detail/cagra/cagra_search.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_plan.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_single_cta.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 2 +- cpp/src/neighbors/detail/cagra/topk_by_radix.cuh | 2 +- cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 53f54a0429..2d383a2429 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 44aa4c9601..c8b885dffe 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/search_plan.cuh b/cpp/src/neighbors/detail/cagra/search_plan.cuh index f7035a23bb..9cc6aeb353 100644 --- a/cpp/src/neighbors/detail/cagra/search_plan.cuh +++ b/cpp/src/neighbors/detail/cagra/search_plan.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh index b2b4670964..0fdf0f208b 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index b0f86d8a81..404817e582 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh index 7b0e3dbcad..be38c6f4b1 100644 --- a/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_by_radix.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh index 97ea7dc236..93b78b8177 100644 --- a/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh +++ b/cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once