From 465844cce2d5ebd5b1f69657fb39956d4d2b992a Mon Sep 17 00:00:00 2001 From: mnorris11 Date: Sun, 27 Sep 2026 14:19:07 -0700 Subject: [PATCH 1/6] break a k-selection tie --- .../raft/matrix/detail/select_warpsort.cuh | 17 +++++++- cpp/include/raft/util/bitonic_sort.cuh | 43 ++++++++++++++++++- 2 files changed, 56 insertions(+), 4 deletions(-) diff --git a/cpp/include/raft/matrix/detail/select_warpsort.cuh b/cpp/include/raft/matrix/detail/select_warpsort.cuh index 3341156cea..5db70278c9 100644 --- a/cpp/include/raft/matrix/detail/select_warpsort.cuh +++ b/cpp/include/raft/matrix/detail/select_warpsort.cuh @@ -109,6 +109,16 @@ _RAFT_DEVICE _RAFT_FORCEINLINE auto is_ordered(T left, T right) -> bool if constexpr (!Ascending) { return left > right; } } +/** `is_ordered` with an equal-key tie broken on the index, as `swap_needed` does. */ +template +_RAFT_DEVICE _RAFT_FORCEINLINE auto is_ordered(T left, T right, IdxT left_idx, IdxT right_idx) + -> bool +{ + if (left != right) { return is_ordered(left, right); } + if constexpr (Ascending) { return left_idx < right_idx; } + if constexpr (!Ascending) { return left_idx > right_idx; } +} + } // namespace /** @@ -188,7 +198,7 @@ class warp_sort { for (int i = kMaxArrLen - 1; i >= 0; --i, idx += kWarpWidth) { if (idx < k) { T t = in[idx]; - if (is_ordered(t, val_arr_[i])) { + if (is_ordered(t, val_arr_[i], in_idx[idx], idx_arr_[i])) { val_arr_[i] = t; idx_arr_[i] = in_idx[idx]; } @@ -259,7 +269,10 @@ class warp_sort { for (int i = std::min(kMaxArrLen, PerThreadSizeIn); i > 0; i--) { T& key = val_arr_[kMaxArrLen - i]; T other = keys_in[PerThreadSizeIn - i]; - if (is_ordered(other, key)) { + if (is_ordered(other, + key, + ids_in[PerThreadSizeIn - i], + idx_arr_[kMaxArrLen - i])) { key = other; idx_arr_[kMaxArrLen - i] = ids_in[PerThreadSizeIn - i]; } diff --git a/cpp/include/raft/util/bitonic_sort.cuh b/cpp/include/raft/util/bitonic_sort.cuh index 318b47e0dd..5c05d0d4d5 100644 --- a/cpp/include/raft/util/bitonic_sort.cuh +++ b/cpp/include/raft/util/bitonic_sort.cuh @@ -27,6 +27,26 @@ _RAFT_DEVICE _RAFT_FORCEINLINE void conditional_assign(bool cond, T& ptr, T x) if (cond) { ptr = x; } } +/** First element of a payload pack. k-selection always passes the index first. */ +template +_RAFT_DEVICE _RAFT_FORCEINLINE P0 first_payload(P0 head, Rest...) +{ + return head; +} + +/** Whether `a` must move after `b`, breaking an equal-key tie on the index. */ +template +_RAFT_DEVICE _RAFT_FORCEINLINE bool swap_needed(bool ascending, + KeyT a, + KeyT b, + IdxT ia, + IdxT ib) +{ + if (a != b) { return ascending ? (a > b) : (a < b); } + // Must flip with `ascending`, or paired lanes disagree on a tie and duplicate an element. + return ascending ? (ia > ib) : (ia < ib); +} + } // namespace /** @@ -192,7 +212,17 @@ class bitonic { const int other_i = i + stride; KeyT& key = keys[i]; KeyT& other = keys[other_i]; - if (ascending ? key > other : key < other) { + bool do_swap; + if constexpr (sizeof...(PayloadTs) > 0) { + do_swap = swap_needed(ascending, + key, + other, + first_payload(payloads[i]...), + first_payload(payloads[other_i]...)); + } else { + do_swap = ascending ? key > other : key < other; + } + if (do_swap) { swap(key, other); (swap(payloads[i], payloads[other_i]), ...); } @@ -206,7 +236,16 @@ class bitonic { for (int stride = (warp_width >> 1); stride > 0; stride >>= 1) { const bool is_second = lane & stride; const KeyT other = shfl_xor(key, stride, warp_width); - const bool do_assign = (ascending != is_second) ? key > other : key < other; + const bool asc = (ascending != is_second); + bool do_assign; + if constexpr (sizeof...(PayloadTs) > 0) { + // The partner's index must cross the lane boundary to break the tie. + const auto my_id = first_payload(payloads[i]...); + const auto other_id = shfl_xor(my_id, stride, warp_width); + do_assign = swap_needed(asc, key, other, my_id, other_id); + } else { + do_assign = asc ? key > other : key < other; + } conditional_assign(do_assign, key, other); // NB: don't put shfl_xor in a conditional; it must be called by all threads in a warp. From 54500b16302b8aefdeb217793df04d314f9f40a0 Mon Sep 17 00:00:00 2001 From: mnorris11 Date: Mon, 28 Sep 2026 15:00:04 -0700 Subject: [PATCH 2/6] Add SelectAlgo::kWarpDistributedShmStable for a reproducible k-selection --- .../raft/matrix/detail/select_k-inl.cuh | 4 ++ .../raft/matrix/detail/select_warpsort.cuh | 57 +++++++++++++------ cpp/include/raft/matrix/select_k_types.hpp | 7 +++ cpp/include/raft/util/bitonic_sort.cuh | 38 ++++++++----- 4 files changed, 74 insertions(+), 32 deletions(-) diff --git a/cpp/include/raft/matrix/detail/select_k-inl.cuh b/cpp/include/raft/matrix/detail/select_k-inl.cuh index bbdbcb886f..3733a478c8 100644 --- a/cpp/include/raft/matrix/detail/select_k-inl.cuh +++ b/cpp/include/raft/matrix/detail/select_k-inl.cuh @@ -299,6 +299,10 @@ void select_k(raft::resources const& handle, return detail::select::warpsort:: select_k_impl( handle, in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min); + case SelectAlgo::kWarpDistributedShmStable: + return detail::select::warpsort:: + select_k_impl( + handle, in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min, nullptr, true); case SelectAlgo::kWarpAuto: return detail::select::warpsort::select_k( handle, in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min); diff --git a/cpp/include/raft/matrix/detail/select_warpsort.cuh b/cpp/include/raft/matrix/detail/select_warpsort.cuh index 5db70278c9..6699c37a00 100644 --- a/cpp/include/raft/matrix/detail/select_warpsort.cuh +++ b/cpp/include/raft/matrix/detail/select_warpsort.cuh @@ -156,6 +156,9 @@ class warp_sort { /** Extra memory required per-block for keeping the state (shared or global). */ constexpr static auto mem_required(uint32_t block_size) -> size_t { return 0; } + /** Order an equal-key tie on the payload index rather than on lane arrival. */ + _RAFT_DEVICE _RAFT_FORCEINLINE void set_tie_break(bool tie_break) { tie_break_ = tie_break; } + /** * Construct the warp_sort empty queue. * @@ -198,7 +201,8 @@ class warp_sort { for (int i = kMaxArrLen - 1; i >= 0; --i, idx += kWarpWidth) { if (idx < k) { T t = in[idx]; - if (is_ordered(t, val_arr_[i], in_idx[idx], idx_arr_[i])) { + if (tie_break_ ? is_ordered(t, val_arr_[i], in_idx[idx], idx_arr_[i]) + : is_ordered(t, val_arr_[i])) { val_arr_[i] = t; idx_arr_[i] = in_idx[idx]; } @@ -206,7 +210,7 @@ class warp_sort { } } if (kWarpWidth < WarpSize || do_merge) { - util::bitonic(Ascending, kWarpWidth).merge(val_arr_, idx_arr_); + util::bitonic(Ascending, kWarpWidth, this->tie_break_).merge(val_arr_, idx_arr_); } } @@ -244,6 +248,8 @@ class warp_sort { T val_arr_[kMaxArrLen]; IdxT idx_arr_[kMaxArrLen]; + /** Order an equal-key tie on the index. Off by default, so the order is unchanged. */ + bool tie_break_ = false; /** * Merge another array (sorted in the opposite direction) in the queue. @@ -269,15 +275,16 @@ class warp_sort { for (int i = std::min(kMaxArrLen, PerThreadSizeIn); i > 0; i--) { T& key = val_arr_[kMaxArrLen - i]; T other = keys_in[PerThreadSizeIn - i]; - if (is_ordered(other, - key, - ids_in[PerThreadSizeIn - i], - idx_arr_[kMaxArrLen - i])) { + if (tie_break_ ? is_ordered(other, + key, + ids_in[PerThreadSizeIn - i], + idx_arr_[kMaxArrLen - i]) + : is_ordered(other, key)) { key = other; idx_arr_[kMaxArrLen - i] = ids_in[PerThreadSizeIn - i]; } } - util::bitonic(Ascending, kWarpWidth).merge(val_arr_, idx_arr_); + util::bitonic(Ascending, kWarpWidth, this->tie_break_).merge(val_arr_, idx_arr_); } }; @@ -347,7 +354,7 @@ class warp_sort_filtered : public warp_sort { _RAFT_DEVICE _RAFT_FORCEINLINE void merge_buf_() { - util::bitonic(!Ascending, kWarpWidth).sort(val_buf_, idx_buf_); + util::bitonic(!Ascending, kWarpWidth, this->tie_break_).sort(val_buf_, idx_buf_); this->template merge_in(val_buf_, idx_buf_); buf_len_ = 0; set_k_th_(); // contains warp sync @@ -472,7 +479,7 @@ class warp_sort_distributed : public warp_sort { _RAFT_DEVICE _RAFT_FORCEINLINE void merge_buf_() { - util::bitonic<1>(!Ascending, kWarpWidth).sort(buf_val_, buf_idx_); + util::bitonic<1>(!Ascending, kWarpWidth, this->tie_break_).sort(buf_val_, buf_idx_); this->template merge_in<1>(&buf_val_, &buf_idx_); set_k_th_(); // contains warp sync buf_val_ = kDummy; @@ -584,7 +591,7 @@ class warp_sort_distributed_ext : public warp_sort T buf_val = val_buf_[laneId()]; IdxT buf_idx = idx_buf_[laneId()]; val_buf_[laneId()] = kDummy; - util::bitonic<1>(!Ascending, kWarpWidth).sort(buf_val, buf_idx); + util::bitonic<1>(!Ascending, kWarpWidth, this->tie_break_).sort(buf_val, buf_idx); this->template merge_in<1>(&buf_val, &buf_idx); set_k_th_(); // contains warp sync } @@ -643,7 +650,7 @@ class warp_sort_immediate : public warp_sort { ++buf_len_; if (buf_len_ == kMaxArrLen) { - util::bitonic(!Ascending, kWarpWidth).sort(val_buf_, idx_buf_); + util::bitonic(!Ascending, kWarpWidth, this->tie_break_).sort(val_buf_, idx_buf_); this->template merge_in(val_buf_, idx_buf_); #pragma unroll for (int i = 0; i < kMaxArrLen; i++) { @@ -656,7 +663,7 @@ class warp_sort_immediate : public warp_sort { _RAFT_DEVICE void done() { if (buf_len_ != 0) { - util::bitonic(!Ascending, kWarpWidth).sort(val_buf_, idx_buf_); + util::bitonic(!Ascending, kWarpWidth, this->tie_break_).sort(val_buf_, idx_buf_); this->template merge_in(val_buf_, idx_buf_); } } @@ -692,6 +699,8 @@ class block_sort { { } + _RAFT_DEVICE void set_tie_break(bool tie_break) { queue_.set_tie_break(tie_break); } + _RAFT_DEVICE void add(T val, IdxT idx) { queue_.add(val, idx); } /** @@ -771,7 +780,8 @@ __launch_bounds__(256) RAFT_KERNEL block_kernel(const T* in, IdxT len, int k, T* out, - IdxT* out_idx) + IdxT* out_idx, + bool tie_break) { // * per-block output { @@ -801,6 +811,7 @@ __launch_bounds__(256) RAFT_KERNEL block_kernel(const T* in, using bq_t = block_sort; uint8_t* warp_smem = bq_t::queue_t::mem_required(blockDim.x) > 0 ? smem_buf_bytes : nullptr; bq_t queue(k, warp_smem); + queue.set_tie_break(tie_break); // * main loop const IdxT stride = gridDim.x * blockDim.x; @@ -878,6 +889,7 @@ struct launch_setup { const IdxT* in_indptr, T* out_key, IdxT* out_idx, + bool tie_break, cuda::stream_ref stream) { const int capacity = bound_by_power_of_two(k); @@ -895,6 +907,7 @@ struct launch_setup { in_indptr, out_key, out_idx, + tie_break, stream); } } @@ -919,7 +932,8 @@ struct launch_setup { IdxT(len), k, out_key, - out_idx); + out_idx, + tie_break); } else { raft::launch_kernel({stream, smem_size}, gs, @@ -932,7 +946,8 @@ struct launch_setup { IdxT(len), k, out_key, - out_idx); + out_idx, + tie_break); } RAFT_CUDA_TRY(cudaPeekAtLastError()); out_key += batch_chunk * num_blocks * k; @@ -1088,6 +1103,7 @@ void select_k_(bool dry_run, T* out, IdxT* out_idx, bool select_min, + bool tie_break, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { @@ -1118,6 +1134,7 @@ void select_k_(bool dry_run, in_indptr, result_val, result_idx, + tie_break, stream); if (num_of_block > 1) { @@ -1134,6 +1151,7 @@ void select_k_(bool dry_run, nullptr, out, out_idx, + tie_break, stream); } } @@ -1151,7 +1169,8 @@ void select_k_impl(raft::resources const& res, T* out, IdxT* out_idx, bool select_min, - const IdxT* in_indptr = nullptr) + const IdxT* in_indptr = nullptr, + bool tie_break = false) { int num_of_block = 0; int num_of_warp = 0; @@ -1170,6 +1189,7 @@ void select_k_impl(raft::resources const& res, out, out_idx, select_min, + tie_break, resource::get_cuda_stream(res), resource::get_workspace_resource_ref(res)); } @@ -1221,7 +1241,8 @@ void select_k(raft::resources const& res, T* out, IdxT* out_idx, bool select_min, - const IdxT* in_indptr = nullptr) + const IdxT* in_indptr = nullptr, + bool tie_break = false) { ASSERT(k <= kMaxCapacity, "Current max k is %d (requested %d)", kMaxCapacity, k); ASSERT(len <= size_t(std::numeric_limits::max()), @@ -1248,6 +1269,7 @@ void select_k(raft::resources const& res, out, out_idx, select_min, + tie_break, resource::get_cuda_stream(res), resource::get_workspace_resource_ref(res)); } else { @@ -1265,6 +1287,7 @@ void select_k(raft::resources const& res, out, out_idx, select_min, + tie_break, resource::get_cuda_stream(res), resource::get_workspace_resource_ref(res)); } diff --git a/cpp/include/raft/matrix/select_k_types.hpp b/cpp/include/raft/matrix/select_k_types.hpp index c6f42468e4..033739a74a 100644 --- a/cpp/include/raft/matrix/select_k_types.hpp +++ b/cpp/include/raft/matrix/select_k_types.hpp @@ -67,6 +67,11 @@ enum class SelectAlgo : uint8_t { * in the given external pointers (normally, a shared memory pointer should be passed in). */ kWarpDistributedShm = 8, + /** + * The same as `kWarpDistributedShm`, but orders an equal-key tie by the payload index instead of + * by lane arrival, so a repeated selection over the same input returns the same elements. + */ + kWarpDistributedShmStable = 9, }; inline auto operator<<(std::ostream& os, const SelectAlgo& algo) -> std::ostream& @@ -74,6 +79,8 @@ inline auto operator<<(std::ostream& os, const SelectAlgo& algo) -> std::ostream auto underlying_value = static_cast::type>(algo); switch (algo) { + case SelectAlgo::kWarpDistributedShmStable: + return os << "kWarpDistributedShmStable=" << underlying_value; case SelectAlgo::kAuto: return os << "kAuto=" << underlying_value; case SelectAlgo::kRadix8bits: return os << "kRadix8bits=" << underlying_value; case SelectAlgo::kRadix11bits: return os << "kRadix11bits=" << underlying_value; diff --git a/cpp/include/raft/util/bitonic_sort.cuh b/cpp/include/raft/util/bitonic_sort.cuh index 5c05d0d4d5..faff3c8bbb 100644 --- a/cpp/include/raft/util/bitonic_sort.cuh +++ b/cpp/include/raft/util/bitonic_sort.cuh @@ -104,8 +104,10 @@ class bitonic { * the total size of the sorted data is `Size * warp_width`. * Must be power-of-two, not larger than the WarpSize. */ - _RAFT_DEVICE _RAFT_FORCEINLINE explicit bitonic(bool ascending, int warp_width = WarpSize) - : ascending_(ascending), warp_width_(warp_width) + _RAFT_DEVICE _RAFT_FORCEINLINE explicit bitonic(bool ascending, + int warp_width = WarpSize, + bool tie_break = false) + : ascending_(ascending), warp_width_(warp_width), tie_break_(tie_break) { } @@ -135,7 +137,7 @@ class bitonic { _RAFT_DEVICE _RAFT_FORCEINLINE void merge(KeyT* __restrict__ keys, PayloadTs* __restrict__... payloads) const { - return bitonic::merge_impl(ascending_, warp_width_, keys, payloads...); + return bitonic::merge_impl(ascending_, warp_width_, tie_break_, keys, payloads...); } /** @@ -154,7 +156,7 @@ class bitonic { _RAFT_DEVICE _RAFT_FORCEINLINE void sort(KeyT* __restrict__ keys, PayloadTs* __restrict__... payloads) const { - return bitonic::sort_impl(ascending_, warp_width_, keys, payloads...); + return bitonic::sort_impl(ascending_, warp_width_, tie_break_, keys, payloads...); } /** @@ -192,6 +194,8 @@ class bitonic { private: const int warp_width_; const bool ascending_; + /** Order an equal-key tie on the payload index. Off by default. */ + const bool tie_break_; template friend class bitonic; @@ -199,6 +203,7 @@ class bitonic { template static _RAFT_DEVICE _RAFT_FORCEINLINE void merge_impl(bool ascending, int warp_width, + bool tie_break, KeyT* __restrict__ keys, PayloadTs* __restrict__... payloads) { @@ -214,11 +219,12 @@ class bitonic { KeyT& other = keys[other_i]; bool do_swap; if constexpr (sizeof...(PayloadTs) > 0) { - do_swap = swap_needed(ascending, - key, - other, - first_payload(payloads[i]...), - first_payload(payloads[other_i]...)); + do_swap = tie_break ? swap_needed(ascending, + key, + other, + first_payload(payloads[i]...), + first_payload(payloads[other_i]...)) + : (ascending ? key > other : key < other); } else { do_swap = ascending ? key > other : key < other; } @@ -239,10 +245,11 @@ class bitonic { const bool asc = (ascending != is_second); bool do_assign; if constexpr (sizeof...(PayloadTs) > 0) { - // The partner's index must cross the lane boundary to break the tie. + // The shuffle is unconditional: it is a warp-collective operation. const auto my_id = first_payload(payloads[i]...); const auto other_id = shfl_xor(my_id, stride, warp_width); - do_assign = swap_needed(asc, key, other, my_id, other_id); + do_assign = tie_break ? swap_needed(asc, key, other, my_id, other_id) + : (asc ? key > other : key < other); } else { do_assign = asc ? key > other : key < other; } @@ -258,20 +265,21 @@ class bitonic { template static _RAFT_DEVICE _RAFT_FORCEINLINE void sort_impl(bool ascending, int warp_width, + bool tie_break, KeyT* __restrict__ keys, PayloadTs* __restrict__... payloads) { if constexpr (Size == 1) { const int lane = laneId(); for (int width = 2; width < warp_width; width <<= 1) { - bitonic<1>::merge_impl(lane & width, width, keys, payloads...); + bitonic<1>::merge_impl(lane & width, width, tie_break, keys, payloads...); } } else { constexpr int kSize2 = Size / 2; - bitonic::sort_impl(false, warp_width, keys, payloads...); - bitonic::sort_impl(true, warp_width, keys + kSize2, (payloads + kSize2)...); + bitonic::sort_impl(false, warp_width, tie_break, keys, payloads...); + bitonic::sort_impl(true, warp_width, tie_break, keys + kSize2, (payloads + kSize2)...); } - bitonic::merge_impl(ascending, warp_width, keys, payloads...); + bitonic::merge_impl(ascending, warp_width, tie_break, keys, payloads...); } }; From 651d9c15775fe6bffc183a1d3a8376665e362411 Mon Sep 17 00:00:00 2001 From: mnorris11 Date: Fri, 2 Oct 2026 09:26:05 -0700 Subject: [PATCH 3/6] Make the k-selection tie break a compile-time choice The tie break was a run-time flag. Every warp-sort kernel then held both comparisons in its unrolled bitonic networks, and shuffled the index on each step even with the flag off. The kernels of every select_k caller grew, and the default path did extra work. - `bitonic` and `warp_sort` take a `TieBreak` template parameter, false by default. With false, the generated code is the same as before the change. - `warp_sort_distributed_ext` takes the same parameter. The new class `warp_sort_distributed_ext_stable` sets it to true. - `SelectAlgo::kWarpDistributedShmStable` uses that class, for float and half keys only. Other key types fail with an error, so their instantiations do not grow. - The `tie_break` arguments of the kernels and of `select_k_impl` are removed. --- .../raft/matrix/detail/select_k-inl.cuh | 13 +- .../raft/matrix/detail/select_warpsort.cuh | 121 ++++++++++-------- cpp/include/raft/util/bitonic_sort.cuh | 49 ++++--- 3 files changed, 103 insertions(+), 80 deletions(-) diff --git a/cpp/include/raft/matrix/detail/select_k-inl.cuh b/cpp/include/raft/matrix/detail/select_k-inl.cuh index 3733a478c8..b0dddb0a3d 100644 --- a/cpp/include/raft/matrix/detail/select_k-inl.cuh +++ b/cpp/include/raft/matrix/detail/select_k-inl.cuh @@ -300,9 +300,16 @@ void select_k(raft::resources const& handle, select_k_impl( handle, in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min); case SelectAlgo::kWarpDistributedShmStable: - return detail::select::warpsort:: - select_k_impl( - handle, in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min, nullptr, true); + // Only the keys that the IVF searches select on, so other instantiations do not grow. + if constexpr (std::is_same_v || std::is_same_v) { + return detail::select::warpsort::select_k_impl< + T, + IdxT, + detail::select::warpsort::warp_sort_distributed_ext_stable, + row_layout>(handle, in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min); + } else { + RAFT_FAIL("kWarpDistributedShmStable supports float and half keys only"); + } case SelectAlgo::kWarpAuto: return detail::select::warpsort::select_k( handle, in_val, in_idx, batch_size, len, k, out_val, out_idx, select_min); diff --git a/cpp/include/raft/matrix/detail/select_warpsort.cuh b/cpp/include/raft/matrix/detail/select_warpsort.cuh index 6699c37a00..36f19fd264 100644 --- a/cpp/include/raft/matrix/detail/select_warpsort.cuh +++ b/cpp/include/raft/matrix/detail/select_warpsort.cuh @@ -136,8 +136,12 @@ _RAFT_DEVICE _RAFT_FORCEINLINE auto is_ordered(T left, T right, IdxT left_idx, I * @tparam IdxT * the type of payload (normally, indices of elements), i.e. * the content sorted alongside the keys. + * @tparam TieBreak + * order an equal-key tie on the index rather than on arrival, so that a repeated selection + * over the same input returns the same elements. Without it the code is the same as before + * the parameter existed. */ -template +template class warp_sort { static_assert(is_a_power_of_two(Capacity)); static_assert(std::is_default_constructible_v); @@ -156,9 +160,6 @@ class warp_sort { /** Extra memory required per-block for keeping the state (shared or global). */ constexpr static auto mem_required(uint32_t block_size) -> size_t { return 0; } - /** Order an equal-key tie on the payload index rather than on lane arrival. */ - _RAFT_DEVICE _RAFT_FORCEINLINE void set_tie_break(bool tie_break) { tie_break_ = tie_break; } - /** * Construct the warp_sort empty queue. * @@ -201,8 +202,13 @@ class warp_sort { for (int i = kMaxArrLen - 1; i >= 0; --i, idx += kWarpWidth) { if (idx < k) { T t = in[idx]; - if (tie_break_ ? is_ordered(t, val_arr_[i], in_idx[idx], idx_arr_[i]) - : is_ordered(t, val_arr_[i])) { + bool take; + if constexpr (TieBreak) { + take = is_ordered(t, val_arr_[i], in_idx[idx], idx_arr_[i]); + } else { + take = is_ordered(t, val_arr_[i]); + } + if (take) { val_arr_[i] = t; idx_arr_[i] = in_idx[idx]; } @@ -210,7 +216,7 @@ class warp_sort { } } if (kWarpWidth < WarpSize || do_merge) { - util::bitonic(Ascending, kWarpWidth, this->tie_break_).merge(val_arr_, idx_arr_); + util::bitonic(Ascending, kWarpWidth).merge(val_arr_, idx_arr_); } } @@ -248,8 +254,6 @@ class warp_sort { T val_arr_[kMaxArrLen]; IdxT idx_arr_[kMaxArrLen]; - /** Order an equal-key tie on the index. Off by default, so the order is unchanged. */ - bool tie_break_ = false; /** * Merge another array (sorted in the opposite direction) in the queue. @@ -275,16 +279,19 @@ class warp_sort { for (int i = std::min(kMaxArrLen, PerThreadSizeIn); i > 0; i--) { T& key = val_arr_[kMaxArrLen - i]; T other = keys_in[PerThreadSizeIn - i]; - if (tie_break_ ? is_ordered(other, - key, - ids_in[PerThreadSizeIn - i], - idx_arr_[kMaxArrLen - i]) - : is_ordered(other, key)) { + bool take; + if constexpr (TieBreak) { + take = is_ordered( + other, key, ids_in[PerThreadSizeIn - i], idx_arr_[kMaxArrLen - i]); + } else { + take = is_ordered(other, key); + } + if (take) { key = other; idx_arr_[kMaxArrLen - i] = ids_in[PerThreadSizeIn - i]; } } - util::bitonic(Ascending, kWarpWidth, this->tie_break_).merge(val_arr_, idx_arr_); + util::bitonic(Ascending, kWarpWidth).merge(val_arr_, idx_arr_); } }; @@ -354,7 +361,7 @@ class warp_sort_filtered : public warp_sort { _RAFT_DEVICE _RAFT_FORCEINLINE void merge_buf_() { - util::bitonic(!Ascending, kWarpWidth, this->tie_break_).sort(val_buf_, idx_buf_); + util::bitonic(!Ascending, kWarpWidth).sort(val_buf_, idx_buf_); this->template merge_in(val_buf_, idx_buf_); buf_len_ = 0; set_k_th_(); // contains warp sync @@ -479,7 +486,7 @@ class warp_sort_distributed : public warp_sort { _RAFT_DEVICE _RAFT_FORCEINLINE void merge_buf_() { - util::bitonic<1>(!Ascending, kWarpWidth, this->tie_break_).sort(buf_val_, buf_idx_); + util::bitonic<1>(!Ascending, kWarpWidth).sort(buf_val_, buf_idx_); this->template merge_in<1>(&buf_val_, &buf_idx_); set_k_th_(); // contains warp sync buf_val_ = kDummy; @@ -500,12 +507,12 @@ class warp_sort_distributed : public warp_sort { * The same as `warp_sort_distributed`, but keeps the temporary value and index buffers * in the given external pointers (normally, a shared memory pointer should be passed in). */ -template -class warp_sort_distributed_ext : public warp_sort { +template +class warp_sort_distributed_ext : public warp_sort { public: - using warp_sort::kDummy; - using warp_sort::kWarpWidth; - using warp_sort::k; + using warp_sort::kDummy; + using warp_sort::kWarpWidth; + using warp_sort::k; constexpr static auto mem_required(uint32_t block_size) -> size_t { @@ -513,7 +520,7 @@ class warp_sort_distributed_ext : public warp_sort } _RAFT_DEVICE warp_sort_distributed_ext(int k, T* val_buf, IdxT* idx_buf, T limit = kDummy) - : warp_sort(k), + : warp_sort(k), val_buf_(val_buf), idx_buf_(idx_buf), buf_len_(0), @@ -536,7 +543,8 @@ class warp_sort_distributed_ext : public warp_sort auto warp_offset = Pow2::roundDown(threadIdx.x); val_buf += warp_offset; idx_buf += warp_offset; - return warp_sort_distributed_ext{k, val_buf, idx_buf, limit}; + return warp_sort_distributed_ext{ + k, val_buf, idx_buf, limit}; } _RAFT_DEVICE void add(T val, IdxT idx) @@ -591,14 +599,14 @@ class warp_sort_distributed_ext : public warp_sort T buf_val = val_buf_[laneId()]; IdxT buf_idx = idx_buf_[laneId()]; val_buf_[laneId()] = kDummy; - util::bitonic<1>(!Ascending, kWarpWidth, this->tie_break_).sort(buf_val, buf_idx); + util::bitonic<1, TieBreak>(!Ascending, kWarpWidth).sort(buf_val, buf_idx); this->template merge_in<1>(&buf_val, &buf_idx); set_k_th_(); // contains warp sync } - using warp_sort::kMaxArrLen; - using warp_sort::val_arr_; - using warp_sort::idx_arr_; + using warp_sort::kMaxArrLen; + using warp_sort::val_arr_; + using warp_sort::idx_arr_; T* val_buf_; IdxT* idx_buf_; @@ -607,6 +615,27 @@ class warp_sort_distributed_ext : public warp_sort T k_th_; }; +/** + * `warp_sort_distributed_ext` that orders an equal-key tie on the index, so a repeated selection + * over the same input returns the same elements. Selected by `SelectAlgo::kWarpDistributedShmStable`. + */ +template +class warp_sort_distributed_ext_stable + : public warp_sort_distributed_ext { + using base_t = warp_sort_distributed_ext; + + public: + using base_t::base_t; + + _RAFT_DEVICE static auto init_blockwide(int k, uint8_t* shmem, T limit = base_t::kDummy) + { + return warp_sort_distributed_ext_stable{base_t::init_blockwide(k, shmem, limit)}; + } + + private: + _RAFT_DEVICE explicit warp_sort_distributed_ext_stable(const base_t& queue) : base_t(queue) {} +}; + /** * This version of warp_sort adds every input element into the intermediate sorting * buffer, and thus does the sorting step every `Capacity` input elements. @@ -650,7 +679,7 @@ class warp_sort_immediate : public warp_sort { ++buf_len_; if (buf_len_ == kMaxArrLen) { - util::bitonic(!Ascending, kWarpWidth, this->tie_break_).sort(val_buf_, idx_buf_); + util::bitonic(!Ascending, kWarpWidth).sort(val_buf_, idx_buf_); this->template merge_in(val_buf_, idx_buf_); #pragma unroll for (int i = 0; i < kMaxArrLen; i++) { @@ -663,7 +692,7 @@ class warp_sort_immediate : public warp_sort { _RAFT_DEVICE void done() { if (buf_len_ != 0) { - util::bitonic(!Ascending, kWarpWidth, this->tie_break_).sort(val_buf_, idx_buf_); + util::bitonic(!Ascending, kWarpWidth).sort(val_buf_, idx_buf_); this->template merge_in(val_buf_, idx_buf_); } } @@ -699,8 +728,6 @@ class block_sort { { } - _RAFT_DEVICE void set_tie_break(bool tie_break) { queue_.set_tie_break(tie_break); } - _RAFT_DEVICE void add(T val, IdxT idx) { queue_.add(val, idx); } /** @@ -780,8 +807,7 @@ __launch_bounds__(256) RAFT_KERNEL block_kernel(const T* in, IdxT len, int k, T* out, - IdxT* out_idx, - bool tie_break) + IdxT* out_idx) { // * per-block output { @@ -811,7 +837,6 @@ __launch_bounds__(256) RAFT_KERNEL block_kernel(const T* in, using bq_t = block_sort; uint8_t* warp_smem = bq_t::queue_t::mem_required(blockDim.x) > 0 ? smem_buf_bytes : nullptr; bq_t queue(k, warp_smem); - queue.set_tie_break(tie_break); // * main loop const IdxT stride = gridDim.x * blockDim.x; @@ -889,7 +914,6 @@ struct launch_setup { const IdxT* in_indptr, T* out_key, IdxT* out_idx, - bool tie_break, cuda::stream_ref stream) { const int capacity = bound_by_power_of_two(k); @@ -907,7 +931,6 @@ struct launch_setup { in_indptr, out_key, out_idx, - tie_break, stream); } } @@ -932,8 +955,7 @@ struct launch_setup { IdxT(len), k, out_key, - out_idx, - tie_break); + out_idx); } else { raft::launch_kernel({stream, smem_size}, gs, @@ -946,8 +968,7 @@ struct launch_setup { IdxT(len), k, out_key, - out_idx, - tie_break); + out_idx); } RAFT_CUDA_TRY(cudaPeekAtLastError()); out_key += batch_chunk * num_blocks * k; @@ -1002,6 +1023,12 @@ struct LaunchThreshold { static constexpr int len_factor_for_single_block = 32; }; +template <> +struct LaunchThreshold { + static constexpr int len_factor_for_multi_block = 2; + static constexpr int len_factor_for_single_block = 32; +}; + template <> struct LaunchThreshold { static constexpr int len_factor_for_choosing = 4; @@ -1103,7 +1130,6 @@ void select_k_(bool dry_run, T* out, IdxT* out_idx, bool select_min, - bool tie_break, cuda::stream_ref stream, rmm::device_async_resource_ref mr) { @@ -1134,7 +1160,6 @@ void select_k_(bool dry_run, in_indptr, result_val, result_idx, - tie_break, stream); if (num_of_block > 1) { @@ -1151,7 +1176,6 @@ void select_k_(bool dry_run, nullptr, out, out_idx, - tie_break, stream); } } @@ -1169,8 +1193,7 @@ void select_k_impl(raft::resources const& res, T* out, IdxT* out_idx, bool select_min, - const IdxT* in_indptr = nullptr, - bool tie_break = false) + const IdxT* in_indptr = nullptr) { int num_of_block = 0; int num_of_warp = 0; @@ -1189,7 +1212,6 @@ void select_k_impl(raft::resources const& res, out, out_idx, select_min, - tie_break, resource::get_cuda_stream(res), resource::get_workspace_resource_ref(res)); } @@ -1241,8 +1263,7 @@ void select_k(raft::resources const& res, T* out, IdxT* out_idx, bool select_min, - const IdxT* in_indptr = nullptr, - bool tie_break = false) + const IdxT* in_indptr = nullptr) { ASSERT(k <= kMaxCapacity, "Current max k is %d (requested %d)", kMaxCapacity, k); ASSERT(len <= size_t(std::numeric_limits::max()), @@ -1269,7 +1290,6 @@ void select_k(raft::resources const& res, out, out_idx, select_min, - tie_break, resource::get_cuda_stream(res), resource::get_workspace_resource_ref(res)); } else { @@ -1287,7 +1307,6 @@ void select_k(raft::resources const& res, out, out_idx, select_min, - tie_break, resource::get_cuda_stream(res), resource::get_workspace_resource_ref(res)); } diff --git a/cpp/include/raft/util/bitonic_sort.cuh b/cpp/include/raft/util/bitonic_sort.cuh index faff3c8bbb..6d2be7d1e9 100644 --- a/cpp/include/raft/util/bitonic_sort.cuh +++ b/cpp/include/raft/util/bitonic_sort.cuh @@ -87,9 +87,13 @@ _RAFT_DEVICE _RAFT_FORCEINLINE bool swap_needed(bool ascending, * number of elements processed in each thread; * i.e. the total data size is `Size * warp_width`. * Must be power-of-two. + * @tparam TieBreak + * order an equal-key tie on the first payload, which k-selection uses for the index, so + * that a repeated sort of the same input gives the same order. Without it the code is the + * same as before the parameter existed. * */ -template +template class bitonic { static_assert(isPo2(Size)); @@ -104,10 +108,8 @@ class bitonic { * the total size of the sorted data is `Size * warp_width`. * Must be power-of-two, not larger than the WarpSize. */ - _RAFT_DEVICE _RAFT_FORCEINLINE explicit bitonic(bool ascending, - int warp_width = WarpSize, - bool tie_break = false) - : ascending_(ascending), warp_width_(warp_width), tie_break_(tie_break) + _RAFT_DEVICE _RAFT_FORCEINLINE explicit bitonic(bool ascending, int warp_width = WarpSize) + : ascending_(ascending), warp_width_(warp_width) { } @@ -137,7 +139,7 @@ class bitonic { _RAFT_DEVICE _RAFT_FORCEINLINE void merge(KeyT* __restrict__ keys, PayloadTs* __restrict__... payloads) const { - return bitonic::merge_impl(ascending_, warp_width_, tie_break_, keys, payloads...); + return bitonic::merge_impl(ascending_, warp_width_, keys, payloads...); } /** @@ -156,7 +158,7 @@ class bitonic { _RAFT_DEVICE _RAFT_FORCEINLINE void sort(KeyT* __restrict__ keys, PayloadTs* __restrict__... payloads) const { - return bitonic::sort_impl(ascending_, warp_width_, tie_break_, keys, payloads...); + return bitonic::sort_impl(ascending_, warp_width_, keys, payloads...); } /** @@ -194,16 +196,13 @@ class bitonic { private: const int warp_width_; const bool ascending_; - /** Order an equal-key tie on the payload index. Off by default. */ - const bool tie_break_; - template + template friend class bitonic; template static _RAFT_DEVICE _RAFT_FORCEINLINE void merge_impl(bool ascending, int warp_width, - bool tie_break, KeyT* __restrict__ keys, PayloadTs* __restrict__... payloads) { @@ -218,13 +217,12 @@ class bitonic { KeyT& key = keys[i]; KeyT& other = keys[other_i]; bool do_swap; - if constexpr (sizeof...(PayloadTs) > 0) { - do_swap = tie_break ? swap_needed(ascending, - key, - other, - first_payload(payloads[i]...), - first_payload(payloads[other_i]...)) - : (ascending ? key > other : key < other); + if constexpr (TieBreak && sizeof...(PayloadTs) > 0) { + do_swap = swap_needed(ascending, + key, + other, + first_payload(payloads[i]...), + first_payload(payloads[other_i]...)); } else { do_swap = ascending ? key > other : key < other; } @@ -244,12 +242,11 @@ class bitonic { const KeyT other = shfl_xor(key, stride, warp_width); const bool asc = (ascending != is_second); bool do_assign; - if constexpr (sizeof...(PayloadTs) > 0) { + if constexpr (TieBreak && sizeof...(PayloadTs) > 0) { // The shuffle is unconditional: it is a warp-collective operation. const auto my_id = first_payload(payloads[i]...); const auto other_id = shfl_xor(my_id, stride, warp_width); - do_assign = tie_break ? swap_needed(asc, key, other, my_id, other_id) - : (asc ? key > other : key < other); + do_assign = swap_needed(asc, key, other, my_id, other_id); } else { do_assign = asc ? key > other : key < other; } @@ -265,21 +262,21 @@ class bitonic { template static _RAFT_DEVICE _RAFT_FORCEINLINE void sort_impl(bool ascending, int warp_width, - bool tie_break, KeyT* __restrict__ keys, PayloadTs* __restrict__... payloads) { if constexpr (Size == 1) { const int lane = laneId(); for (int width = 2; width < warp_width; width <<= 1) { - bitonic<1>::merge_impl(lane & width, width, tie_break, keys, payloads...); + bitonic<1, TieBreak>::merge_impl(lane & width, width, keys, payloads...); } } else { constexpr int kSize2 = Size / 2; - bitonic::sort_impl(false, warp_width, tie_break, keys, payloads...); - bitonic::sort_impl(true, warp_width, tie_break, keys + kSize2, (payloads + kSize2)...); + bitonic::sort_impl(false, warp_width, keys, payloads...); + bitonic::sort_impl( + true, warp_width, keys + kSize2, (payloads + kSize2)...); } - bitonic::merge_impl(ascending, warp_width, tie_break, keys, payloads...); + bitonic::merge_impl(ascending, warp_width, keys, payloads...); } }; From 4304218e8233098e14c6e3e35df4c7313fc22c18 Mon Sep 17 00:00:00 2001 From: mnorris11 Date: Wed, 7 Oct 2026 16:33:36 -0700 Subject: [PATCH 4/6] Break a tie on the index in the stable k-selection filter `SelectAlgo::kWarpDistributedShmStable` must return the same elements for the same input, whatever the input order. Two defects broke that. `warp_sort_distributed_ext::add` filtered a candidate on its key alone. When the queue was full, it rejected a candidate whose key equalled the k-th key, even when its index should win the tie, so the selected set depended on the input order. With `TieBreak`, the queue now keeps the index of the k-th element, `k_th_idx_`, and the filter compares (key, index) with `is_ordered(val, k_th_, idx, k_th_idx_)`. An empty slot carried index 0, so a real key that twiddles to the same bits as `kDummy` could lose the tie to an empty slot. An empty slot now carries `kDummyIdx`, which loses every tie: the largest `IdxT` for an ascending selection and the smallest for a descending one. The queue start-up, the buffer reset and the padding in `block_kernel` use it. `block_kernel` trims `k` to `len`, so an empty slot never reaches the output. Without `TieBreak` the generated code does not change: every new step is under `if constexpr (TieBreak)`, and `kDummyIdx` is `IdxT{}`. `cpp/tests/matrix/select_k_stable.cu` compares the selector with a CPU sort by key, then by index, for two input orders of the same pairs, and checks that real keys equal to the empty key are selected. --- .../raft/matrix/detail/select_warpsort.cuh | 41 +++-- cpp/tests/CMakeLists.txt | 2 +- cpp/tests/matrix/select_k_stable.cu | 155 ++++++++++++++++++ 3 files changed, 187 insertions(+), 11 deletions(-) create mode 100644 cpp/tests/matrix/select_k_stable.cu diff --git a/cpp/include/raft/matrix/detail/select_warpsort.cuh b/cpp/include/raft/matrix/detail/select_warpsort.cuh index 36f19fd264..7679414a4d 100644 --- a/cpp/include/raft/matrix/detail/select_warpsort.cuh +++ b/cpp/include/raft/matrix/detail/select_warpsort.cuh @@ -137,9 +137,9 @@ _RAFT_DEVICE _RAFT_FORCEINLINE auto is_ordered(T left, T right, IdxT left_idx, I * the type of payload (normally, indices of elements), i.e. * the content sorted alongside the keys. * @tparam TieBreak - * order an equal-key tie on the index rather than on arrival, so that a repeated selection - * over the same input returns the same elements. Without it the code is the same as before - * the parameter existed. + * order an equal-key tie on the index rather than on arrival, so that the selected elements + * do not depend on the input order. Without it the code is the same as before the parameter + * existed. */ template class warp_sort { @@ -152,6 +152,13 @@ class warp_sort { * i.e. `Ascending ? upper_bound() : lower_bound()`. */ static constexpr T kDummy = Ascending ? upper_bound() : lower_bound(); + /** + * The index of an empty slot. With `TieBreak` it loses every tie, so an empty slot never + * outranks a real element whose key equals `kDummy`. + */ + static constexpr IdxT kDummyIdx = + TieBreak ? (Ascending ? upper_bound() : lower_bound()) : IdxT{}; + static constexpr bool kTieBreak = TieBreak; /** Width of the subwarp. */ static constexpr int kWarpWidth = std::min(Capacity, WarpSize); /** The number of elements to select. */ @@ -171,7 +178,7 @@ class warp_sort { #pragma unroll for (int i = 0; i < kMaxArrLen; i++) { val_arr_[i] = kDummy; - idx_arr_[i] = IdxT{}; + idx_arr_[i] = kDummyIdx; } } @@ -511,6 +518,7 @@ template { public: using warp_sort::kDummy; + using warp_sort::kDummyIdx; using warp_sort::kWarpWidth; using warp_sort::k; @@ -524,9 +532,11 @@ class warp_sort_distributed_ext : public warp_sort(val, k_th_); + bool do_add; + if constexpr (TieBreak) { + do_add = is_ordered(val, k_th_, idx, k_th_idx_); + } else { + do_add = is_ordered(val, k_th_); + } // mask tells which lanes in the warp have valid items to be added uint32_t mask = ballot(do_add); if (mask == 0) { return; } @@ -591,6 +606,7 @@ class warp_sort_distributed_ext : public warp_sort(!Ascending, kWarpWidth).sort(buf_val, buf_idx); this->template merge_in<1>(&buf_val, &buf_idx); set_k_th_(); // contains warp sync @@ -613,11 +630,12 @@ class warp_sort_distributed_ext : public warp_sort class warp_sort_distributed_ext_stable @@ -835,6 +853,7 @@ __launch_bounds__(256) RAFT_KERNEL block_kernel(const T* in, // 2. Fewer expensive template instantiations using bits_t = typename cub::Traits::UnsignedBits; using bq_t = block_sort; + using warp_sort_t = WarpSortClass; uint8_t* warp_smem = bq_t::queue_t::mem_required(blockDim.x) > 0 ? smem_buf_bytes : nullptr; bq_t queue(k, warp_smem); @@ -844,8 +863,10 @@ __launch_bounds__(256) RAFT_KERNEL block_kernel(const T* in, for (IdxT i = threadIdx.x + blockIdx.x * blockDim.x; i < per_thread_lim; i += stride) { // Twiddle the input value to ensure proper comparison of floating-point and signed int values queue.add(i < len ? cub::Traits::TwiddleIn(__ldcs(reinterpret_cast(in) + i)) - : WarpSortClass::kDummy, - (i < len && in_idx != nullptr) ? __ldcs(in_idx + i) : i); + : warp_sort_t::kDummy, + (i < len && in_idx != nullptr) ? __ldcs(in_idx + i) + : (warp_sort_t::kTieBreak && i >= len) ? warp_sort_t::kDummyIdx + : i); } // * write out the result diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index e1f7a1f755..36184daf56 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -219,7 +219,7 @@ if(BUILD_TESTS) ) ConfigureTest( - NAME MATRIX_SELECT_TEST PATH matrix/select_k.cu matrix/select_k_edgecases.cu LIB + NAME MATRIX_SELECT_TEST PATH matrix/select_k.cu matrix/select_k_edgecases.cu matrix/select_k_stable.cu LIB EXPLICIT_INSTANTIATE_ONLY GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/matrix/select_k_stable.cu b/cpp/tests/matrix/select_k_stable.cu new file mode 100644 index 0000000000..36445a08a1 --- /dev/null +++ b/cpp/tests/matrix/select_k_stable.cu @@ -0,0 +1,155 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace raft::matrix { + +namespace { + +// The ids that `kWarpDistributedShmStable` must return, sorted: the best keys, and on an equal key +// the smaller id for select-min and the larger id for select-max. +auto expected_ids(const std::vector& keys, + const std::vector& ids, + int k, + bool select_min) -> std::vector +{ + std::vector order(keys.size()); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](size_t a, size_t b) { + if (keys[a] != keys[b]) { return select_min ? keys[a] < keys[b] : keys[a] > keys[b]; } + return select_min ? ids[a] < ids[b] : ids[a] > ids[b]; + }); + std::vector out; + for (int i = 0; i < k; i++) { + out.push_back(ids[order[i]]); + } + std::sort(out.begin(), out.end()); + return out; +} + +auto stable_select(const raft::device_resources& handle, + const std::vector& keys, + const std::vector& ids, + int k, + bool select_min) -> std::vector +{ + auto stream = raft::resource::get_cuda_stream(handle); + int64_t len = keys.size(); + auto in_keys = raft::make_device_matrix(handle, 1, len); + auto in_ids = raft::make_device_matrix(handle, 1, len); + auto out_keys = raft::make_device_matrix(handle, 1, k); + auto out_ids = raft::make_device_matrix(handle, 1, k); + raft::copy(in_keys.data_handle(), keys.data(), len, stream); + raft::copy(in_ids.data_handle(), ids.data(), len, stream); + raft::matrix::select_k(handle, + raft::make_const_mdspan(in_keys.view()), + std::make_optional(raft::make_const_mdspan(in_ids.view())), + out_keys.view(), + out_ids.view(), + select_min, + /*sorted=*/true, + SelectAlgo::kWarpDistributedShmStable); + std::vector out(k); + raft::copy(out.data(), out_ids.data_handle(), k, stream); + raft::resource::sync_stream(handle, stream); + std::sort(out.begin(), out.end()); + return out; +} + +// Selects from the same (key, id) pairs in two input orders. +void check_order_independent(const std::vector& keys, const std::vector& ids, int k) +{ + raft::device_resources handle; + std::mt19937 rng(1234); + for (bool select_min : {true, false}) { + auto expected = expected_ids(keys, ids, k, select_min); + for (int order = 0; order < 2; order++) { + std::vector perm(keys.size()); + std::iota(perm.begin(), perm.end(), 0); + std::shuffle(perm.begin(), perm.end(), rng); + std::vector perm_keys(keys.size()); + std::vector perm_ids(keys.size()); + for (size_t i = 0; i < perm.size(); i++) { + perm_keys[i] = keys[perm[i]]; + perm_ids[i] = ids[perm[i]]; + } + EXPECT_EQ(stable_select(handle, perm_keys, perm_ids, k, select_min), expected) + << "select_min=" << select_min << " order=" << order << " len=" << keys.size() + << " k=" << k; + } + } +} + +auto distinct_ids(size_t n) -> std::vector +{ + std::vector ids(n); + std::iota(ids.begin(), ids.end(), 0); + std::shuffle(ids.begin(), ids.end(), std::mt19937(42)); + return ids; +} + +auto from_bits(uint32_t bits) -> float +{ + float f; + std::memcpy(&f, &bits, sizeof(f)); + return f; +} + +} // namespace + +TEST(SelectKStable, AllKeysEqual) +{ + for (size_t len : {4096, 300000}) { + for (int k : {10, 100}) { + check_order_independent(std::vector(len, 1.0f), distinct_ids(len), k); + } + } +} + +TEST(SelectKStable, TiesAtTheKthPlace) +{ + for (size_t len : {4096, 300000}) { + std::vector keys(len); + std::mt19937 rng(7); + for (auto& key : keys) { + key = float(rng() % 4); + } + check_order_independent(keys, distinct_ids(len), 100); + } +} + +// A real key that twiddles to the same bits as the empty-slot key must outrank the empty slots. +// These NaN bit patterns twiddle to the all-ones (select-min) and all-zeros (select-max) keys. +TEST(SelectKStable, RealKeysEqualToTheEmptyKeyWin) +{ + raft::device_resources handle; + std::vector ids = {11, 3, 7, 5, 9}; + for (bool select_min : {true, false}) { + float key = from_bits(select_min ? 0x7FFFFFFFu : 0xFFFFFFFFu); + std::vector keys(ids.size(), key); + auto got = stable_select(handle, keys, ids, 8, select_min); + for (int64_t id : ids) { + EXPECT_NE(std::find(got.begin(), got.end(), id), got.end()) + << "select_min=" << select_min << " missing id " << id; + } + } +} + +} // namespace raft::matrix From 51f2e0f0141af7858072ab1e1d9919bb995391ef Mon Sep 17 00:00:00 2001 From: mnorris11 Date: Thu, 8 Oct 2026 20:04:35 -0700 Subject: [PATCH 5/6] Store the stable filter's threshold index only with TieBreak Without TieBreak, k_th_idx_ is a cuda::std::monostate marked [[no_unique_address]], so the default selectors keep no extra state. --- cpp/include/raft/matrix/detail/select_warpsort.cuh | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cpp/include/raft/matrix/detail/select_warpsort.cuh b/cpp/include/raft/matrix/detail/select_warpsort.cuh index 7679414a4d..1e0a06d355 100644 --- a/cpp/include/raft/matrix/detail/select_warpsort.cuh +++ b/cpp/include/raft/matrix/detail/select_warpsort.cuh @@ -23,6 +23,7 @@ #include #include // cub::Traits +#include // cuda::std::monostate #include #include @@ -532,11 +533,13 @@ class warp_sort_distributed_ext : public warp_sort k_th_idx_; }; /** From d03e0fa5f0b6a4db58d92eccb85849d56af107a5 Mon Sep 17 00:00:00 2001 From: mnorris11 Date: Wed, 14 Oct 2026 19:26:47 -0700 Subject: [PATCH 6/6] Apply pre-commit formatting and the canonical copyright notice --- cpp/include/raft/matrix/detail/select_warpsort.cuh | 6 +++--- cpp/include/raft/matrix/select_k_types.hpp | 2 +- cpp/include/raft/util/bitonic_sort.cuh | 11 +++-------- cpp/tests/CMakeLists.txt | 4 ++-- cpp/tests/matrix/select_k_stable.cu | 2 +- 5 files changed, 10 insertions(+), 15 deletions(-) diff --git a/cpp/include/raft/matrix/detail/select_warpsort.cuh b/cpp/include/raft/matrix/detail/select_warpsort.cuh index 1e0a06d355..51b6a2bc20 100644 --- a/cpp/include/raft/matrix/detail/select_warpsort.cuh +++ b/cpp/include/raft/matrix/detail/select_warpsort.cuh @@ -23,7 +23,7 @@ #include #include // cub::Traits -#include // cuda::std::monostate +#include // cuda::std::monostate #include #include @@ -289,8 +289,8 @@ class warp_sort { T other = keys_in[PerThreadSizeIn - i]; bool take; if constexpr (TieBreak) { - take = is_ordered( - other, key, ids_in[PerThreadSizeIn - i], idx_arr_[kMaxArrLen - i]); + take = + is_ordered(other, key, ids_in[PerThreadSizeIn - i], idx_arr_[kMaxArrLen - i]); } else { take = is_ordered(other, key); } diff --git a/cpp/include/raft/matrix/select_k_types.hpp b/cpp/include/raft/matrix/select_k_types.hpp index 033739a74a..8c239bb531 100644 --- a/cpp/include/raft/matrix/select_k_types.hpp +++ b/cpp/include/raft/matrix/select_k_types.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/raft/util/bitonic_sort.cuh b/cpp/include/raft/util/bitonic_sort.cuh index 6d2be7d1e9..23aca23bde 100644 --- a/cpp/include/raft/util/bitonic_sort.cuh +++ b/cpp/include/raft/util/bitonic_sort.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -36,11 +36,7 @@ _RAFT_DEVICE _RAFT_FORCEINLINE P0 first_payload(P0 head, Rest...) /** Whether `a` must move after `b`, breaking an equal-key tie on the index. */ template -_RAFT_DEVICE _RAFT_FORCEINLINE bool swap_needed(bool ascending, - KeyT a, - KeyT b, - IdxT ia, - IdxT ib) +_RAFT_DEVICE _RAFT_FORCEINLINE bool swap_needed(bool ascending, KeyT a, KeyT b, IdxT ia, IdxT ib) { if (a != b) { return ascending ? (a > b) : (a < b); } // Must flip with `ascending`, or paired lanes disagree on a tie and duplicate an element. @@ -273,8 +269,7 @@ class bitonic { } else { constexpr int kSize2 = Size / 2; bitonic::sort_impl(false, warp_width, keys, payloads...); - bitonic::sort_impl( - true, warp_width, keys + kSize2, (payloads + kSize2)...); + bitonic::sort_impl(true, warp_width, keys + kSize2, (payloads + kSize2)...); } bitonic::merge_impl(ascending, warp_width, keys, payloads...); } diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 36184daf56..2d12815e1d 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -219,8 +219,8 @@ if(BUILD_TESTS) ) ConfigureTest( - NAME MATRIX_SELECT_TEST PATH matrix/select_k.cu matrix/select_k_edgecases.cu matrix/select_k_stable.cu LIB - EXPLICIT_INSTANTIATE_ONLY GPUS 1 PERCENT 100 + NAME MATRIX_SELECT_TEST PATH matrix/select_k.cu matrix/select_k_edgecases.cu + matrix/select_k_stable.cu LIB EXPLICIT_INSTANTIATE_ONLY GPUS 1 PERCENT 100 ) ConfigureTest( diff --git a/cpp/tests/matrix/select_k_stable.cu b/cpp/tests/matrix/select_k_stable.cu index 36445a08a1..b9e802cdde 100644 --- a/cpp/tests/matrix/select_k_stable.cu +++ b/cpp/tests/matrix/select_k_stable.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */