From 6eeb08f75929b1499afca946bd1a442fffb8b43c Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 9 Jun 2026 01:40:13 -0700 Subject: [PATCH 01/18] exposes batched top-k for small segs --- cub/cub/device/device_batched_topk.cuh | 711 ++++++++++++++++++ .../dispatch/kernels/kernel_batched_topk.cuh | 6 +- .../example_device_batched_topk_keys.cu | 137 ++++ .../example_device_batched_topk_pairs.cu | 158 ++++ .../catch2_test_device_batched_topk_api.cu | 259 +++++++ ...catch2_test_device_batched_topk_env_api.cu | 211 ++++++ .../catch2_test_device_segmented_topk_keys.cu | 52 +- ...catch2_test_device_segmented_topk_pairs.cu | 71 +- 8 files changed, 1563 insertions(+), 42 deletions(-) create mode 100644 cub/cub/device/device_batched_topk.cuh create mode 100644 cub/examples/device/example_device_batched_topk_keys.cu create mode 100644 cub/examples/device/example_device_batched_topk_pairs.cu create mode 100644 cub/test/catch2_test_device_batched_topk_api.cu create mode 100644 cub/test/catch2_test_device_batched_topk_env_api.cu diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh new file mode 100644 index 000000000000..8d798272d609 --- /dev/null +++ b/cub/cub/device/device_batched_topk.cuh @@ -0,0 +1,711 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! @file +//! cub::DeviceBatchedTopK provides device-wide, parallel operations for finding the K largest (or smallest) items +//! from many (small) segments of unordered data items residing within device-accessible memory. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include // topk::select::{min, max} +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +//! @cond +//! Shared implementation for all cub::DeviceBatchedTopK entry points. +//! +//! Validates the requested execution requirements and argument annotations, resolves the (optionally tuned) policy +//! selector from the environment, and forwards to the internal batched top-k dispatch. The selection direction is +//! threaded through as a compile-time `cuda::args::constant` so the kernel only emits the +//! requested (max OR min) code path. +//! +//! All current API-surface constraints are surfaced here as `static_assert`s so the diagnostic appears at the +//! `cub::DeviceBatchedTopK` call site rather than deep inside the kernel/agent instantiation. +template +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + ValueInputIteratorItT d_values_in, + ValueOutputIteratorItT d_values_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env) +{ + // --------------------------------------------------------------------------- + // Execution requirements (mirrors cub::DeviceTopK): non-deterministic + unsorted output only. + // --------------------------------------------------------------------------- + static_assert(!::cuda::std::execution::__queryable_with, + "Determinism should be used inside requires to have an effect."); + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + using requested_order_t = + ::cuda::std::execution::__query_result_or_t; + static_assert(::cuda::std::is_same_v + && ::cuda::std::is_same_v, + "cub::DeviceBatchedTopK only supports non-deterministic, unsorted output. Acknowledge this by " + "passing cuda::execution::require(cuda::execution::determinism::not_guaranteed, " + "cuda::execution::output_ordering::unsorted) in the environment."); + + // A tie-break requirement only constrains *which* of the elements that compare equal at the K-th position are + // selected (per segment); it is meaningless without determinism, so it must be paired with a deterministic + // requirement (run_to_run / gpu_to_gpu). Deterministic execution and tie-breaking are not implemented yet; this only + // validates the requirement combination so the eventual behavior is wired up. + using requested_tie_break_t = + ::cuda::std::execution::__query_result_or_t; + static_assert( + ::cuda::std::is_same_v + || !::cuda::std::is_same_v, + "cub::DeviceBatchedTopK: a tie_break requirement (cuda::execution::tie_break::prefer_smaller_index or " + "prefer_larger_index) requires a deterministic execution requirement " + "(cuda::execution::determinism::run_to_run or gpu_to_gpu); it cannot be combined with " + "cuda::execution::determinism::not_guaranteed."); + + // --------------------------------------------------------------------------- + // Resolve the (optionally tuned) policy selector from the environment. + // --------------------------------------------------------------------------- + using key_t = it_value_t>; + using value_t = it_value_t>; + using default_policy_selector_t = batched_topk:: + policy_selector_from_types::highest>; + using tuning_env_t = + ::cuda::__call_result_or_t<::cuda::execution::__get_tuning_t, ::cuda::std::execution::env<>, EnvT>; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + + // --------------------------------------------------------------------------- + // Argument-annotation constraints surfaced at the call site. + // --------------------------------------------------------------------------- + static_assert(::cuda::args::__traits::is_single_value, + "cub::DeviceBatchedTopK currently requires a single (uniform) number of segments resolved on the " + "host; pass num_segments as a single-value annotation (e.g. cuda::args::constant or " + "cuda::args::immediate), not a per-segment sequence."); + + auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); + + return batched_topk::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + ::cuda::args::constant{}, + num_segments, + total_num_items, + stream.get(), + policy_selector_t{}); +} +//! @endcond +} // namespace detail + +//! @rst +//! DeviceBatchedTopK provides device-wide, parallel operations for finding the largest (or smallest) K items from +//! many segments of unordered data items residing within device-accessible memory. +//! +//! Overview +//! ++++++++++++++++++++++++++ +//! +//! Given a batch of segments, ``DeviceBatchedTopK`` finds, independently for each segment, the K largest (or +//! smallest) items. Within each output segment the K results are returned in an **unspecified order**. The algorithm +//! reuses the block-level top-k primitive used by :cpp:struct:`cub::DeviceTopK`. +//! +//! Argument annotation framework +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The per-segment parameters (segment sizes, ``k``, the number of segments, and a guarantee on the total number of +//! input items) are supplied as **annotated arguments** from ``cuda::args``. This lets callers communicate +//! compile-time guarantees (e.g. a static upper bound on the segment size) that the algorithm uses to specialize the +//! kernel: +//! +//! - ``cuda::args::constant`` — a value known at compile time. +//! - ``cuda::args::immediate{value, cuda::args::bounds()}`` — a runtime value with a +//! compile-time ``[lo, hi]`` range. +//! - ``cuda::args::deferred_sequence{iterator, cuda::args::bounds()}`` — a per-segment +//! sequence (e.g. variable segment sizes) with a compile-time element range. +//! +//! Choosing argument bounds +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! Prefer **sharp (tight) upper bounds**, especially for the segment size. The statically-known *maximum* segment size +//! (the upper bound of the ``segment_sizes`` annotation) does more than select the kernel: it can also drive how much +//! temporary storage the algorithm requests. As a rough intuition, the temporary allocation may grow with the number +//! of segments times some factor of the *maximum* segment size, so an unnecessarily loose upper bound can inflate +//! temporary storage even when the actual segments are much smaller. The precise relationship is intentionally left +//! unspecified and may change across releases (temporary-storage handling is an implementation detail). Treat this +//! purely as guidance for choosing bounds rather than as a guarantee. +//! +//! Current constraints (initial API surface) +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! This is an initial, intentionally restricted API surface. The following constraints are enforced at compile time +//! (a ``static_assert`` fires if violated): +//! +//! - **Small segments only.** Every segment must be processable by a single thread block (one worker per segment). +//! The *statically-known maximum* segment size (the upper bound of the ``segment_sizes`` annotation) must be small +//! enough that such a block fits within the shared-memory limit. Both uniform (fixed) and variable segment sizes are +//! supported as long as this maximum is honored. +//! - **Uniform number of segments.** ``num_segments`` must be a single (host-resolved) value. +//! - **Non-deterministic, unsorted output.** Like :cpp:struct:`cub::DeviceTopK`, the caller must acknowledge this by +//! passing ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, +//! cuda::execution::output_ordering::unsorted)`` in the environment. +//! +//! Determinism +//! ++++++++++++++++++++++++++ +//! +//! DeviceBatchedTopK currently only supports unordered output, which may be non-deterministic for certain inputs: +//! if multiple items tie at the K-th position, the subset of tied elements returned may vary between runs. +//! +//! Usage Considerations +//! ++++++++++++++++++++++++++ +//! +//! @cdp_class{DeviceBatchedTopK} +//! +//! @endrst +struct DeviceBatchedTopK +{ + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Finds, for each segment, the largest K keys from an unordered input sequence of keys. + //! + //! .. note:: + //! + //! The behavior is undefined if the input and output ranges overlap in any way. + //! + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin batched-topk-max-keys + //! :end-before: example-end batched-topk-max-keys + //! + //! @endrst + //! + //! @tparam KeyInputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator + //! + //! @tparam KeyOutputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator + //! + //! @param[in] d_temp_storage + //! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to + //! `temp_storage_bytes` and no work is done. + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] d_keys_in + //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` + //! + //! @param[out] d_keys_out + //! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i` + //! + //! @param[in] segment_sizes + //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum; + //! prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! *Choosing argument bounds* section). + //! + //! @param[in] k + //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! + //! @param[in] num_segments + //! Annotated argument providing the (uniform) number of segments + //! + //! @param[in] total_num_items + //! Annotated argument providing an upper-bound guarantee on the total number of input items across all segments, + //! used to size internal offset types + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed` and + //! `output_ordering::unsorted`. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MaxKeys"); + return detail::dispatch_batched_topk( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + num_segments, + total_num_items, + ::cuda::std::move(env)); + } + + //! @rst + //! Finds, for each segment, the largest K keys from an unordered input sequence of keys. + //! + //! This is an environment-based API that allocates and manages the required temporary storage internally using the + //! memory resource queried from the environment. + //! + //! .. note:: + //! + //! The behavior is undefined if the input and output ranges overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin batched-topk-max-keys-env + //! :end-before: example-end batched-topk-max-keys-env + //! + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys( + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxKeys"); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { + return detail::dispatch_batched_topk( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + num_segments, + total_num_items, + env); + }); + } + + //! @rst + //! Finds, for each segment, the smallest K keys from an unordered input sequence of keys. + //! + //! .. note:: + //! + //! The behavior is undefined if the input and output ranges overlap in any way. + //! + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin batched-topk-min-keys + //! :end-before: example-end batched-topk-min-keys + //! + //! @endrst + //! + //! @copydetails MaxKeys + template > + CUB_RUNTIME_FUNCTION static cudaError_t MinKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MinKeys"); + return detail::dispatch_batched_topk( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + num_segments, + total_num_items, + ::cuda::std::move(env)); + } + + //! @rst + //! Finds, for each segment, the smallest K keys from an unordered input sequence of keys. Environment-based overload + //! that allocates temporary storage internally. + //! + //! .. note:: + //! + //! The behavior is undefined if the input and output ranges overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin batched-topk-min-keys-env + //! :end-before: example-end batched-topk-min-keys-env + //! + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinKeys( + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinKeys"); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { + return detail::dispatch_batched_topk( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + num_segments, + total_num_items, + env); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Finds, for each segment, the largest K keys and their corresponding values from an unordered input sequence of + //! key-value pairs. + //! + //! .. note:: + //! + //! The behavior is undefined if the input and output ranges overlap in any way. + //! + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin batched-topk-max-pairs + //! :end-before: example-end batched-topk-max-pairs + //! + //! @endrst + //! + //! @tparam ValueInputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment value-input iterators @iterator + //! + //! @tparam ValueOutputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment value-output iterators @iterator + //! + //! @param[in] d_values_in + //! Iterator such that `d_values_in[i]` yields a random-access iterator to the values of segment `i` + //! + //! @param[out] d_values_out + //! Iterator such that `d_values_out[i]` yields a random-access output iterator for the values corresponding to the + //! top-k keys of segment `i` + //! + //! @copydetails MaxKeys + template > + CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + ValueInputIteratorItT d_values_in, + ValueOutputIteratorItT d_values_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MaxPairs"); + return detail::dispatch_batched_topk( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + num_segments, + total_num_items, + ::cuda::std::move(env)); + } + + //! @rst + //! Finds, for each segment, the largest K keys and their corresponding values. Environment-based overload that + //! allocates temporary storage internally. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin batched-topk-max-pairs-env + //! :end-before: example-end batched-topk-max-pairs-env + //! + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs( + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + ValueInputIteratorItT d_values_in, + ValueOutputIteratorItT d_values_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxPairs"); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { + return detail::dispatch_batched_topk( + storage, + bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + num_segments, + total_num_items, + env); + }); + } + + //! @rst + //! Finds, for each segment, the smallest K keys and their corresponding values from an unordered input sequence of + //! key-value pairs. + //! + //! .. note:: + //! + //! The behavior is undefined if the input and output ranges overlap in any way. + //! + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin batched-topk-min-pairs + //! :end-before: example-end batched-topk-min-pairs + //! + //! @endrst + //! + //! @copydetails MaxPairs + template > + CUB_RUNTIME_FUNCTION static cudaError_t MinPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + ValueInputIteratorItT d_values_in, + ValueOutputIteratorItT d_values_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MinPairs"); + return detail::dispatch_batched_topk( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + num_segments, + total_num_items, + ::cuda::std::move(env)); + } + + //! @rst + //! Finds, for each segment, the smallest K keys and their corresponding values. Environment-based overload that + //! allocates temporary storage internally. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin batched-topk-min-pairs-env + //! :end-before: example-end batched-topk-min-pairs-env + //! + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinPairs( + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + ValueInputIteratorItT d_values_in, + ValueOutputIteratorItT d_values_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items, + EnvT env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinPairs"); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { + return detail::dispatch_batched_topk( + storage, + bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + num_segments, + total_num_items, + env); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh index 3412f40359a7..9c66be1cba32 100644 --- a/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh +++ b/cub/cub/device/dispatch/kernels/kernel_batched_topk.cuh @@ -81,7 +81,11 @@ private: public: // TODO (elstehle): extend support for variable-size segments - static_assert(selected_index >= 0, "No valid policy found for one-worker-per-segment approach"); + static_assert(selected_index >= 0, + "cub::DeviceBatchedTopK currently supports only segments small enough to be processed by a single " + "thread block (one worker per segment). No policy could cover the statically-known maximum segment " + "size within the shared-memory limit. Reduce the maximum segment size encoded in the segment-size " + "argument annotation."); static constexpr policy_t policy = { active_policy.worker_per_segment_policies[selected_index], active_policy.multi_worker_per_segment_policy}; diff --git a/cub/examples/device/example_device_batched_topk_keys.cu b/cub/examples/device/example_device_batched_topk_keys.cu new file mode 100644 index 000000000000..605552300dc3 --- /dev/null +++ b/cub/examples/device/example_device_batched_topk_keys.cu @@ -0,0 +1,137 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Simple example of cub::DeviceBatchedTopK::MaxKeys(). +//! For each (small) segment, find the K largest float keys. This example also demonstrates the argument annotation +//! framework: the (small) segment size and K are passed as compile-time `cuda::args::constant<>` values, while +//! the number of segments is a runtime `cuda::args::immediate`. +//! +//! To compile using the command line: +//! nvcc -arch=sm_XX example_device_batched_topk_keys.cu -I../.. -lcudart -O3 + +// Ensure printing of CUDA runtime errors to console +#define CUB_STDERR + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "../../test/test_util.h" + +using namespace cub; + +// Whether to display input/output to console +bool g_verbose = false; + +// Compile-time, statically-known (small) segment size and K. Passing them as compile-time constants lets +// cub::DeviceBatchedTopK specialize the kernel for a single thread block per segment. +static constexpr int segment_size = 256; +static constexpr int k = 8; + +int main(int argc, char** argv) +{ + int num_segments = 4000; // runtime number of segments + + // Initialize command line + CommandLineArgs args(argc, argv); + g_verbose = args.CheckCmdLineFlag("v"); + args.GetCmdLineArgument("num-segments", num_segments); + + if (args.CheckCmdLineFlag("help")) + { + std::cout << "Usage: " << argv[0] << " [--num-segments=] [--device=] [--v]\n"; + exit(0); + } + + CubDebugExit(args.DeviceInit()); + + const int num_items = num_segments * segment_size; + std::cout << "cub::DeviceBatchedTopK::MaxKeys() finds the " << k << " largest items in each of " << num_segments + << " segments of " << segment_size << " items (" << sizeof(float) << "-byte keys)\n"; + + // Initialize host input + thrust::host_vector h_keys_in(num_items); + for (int i = 0; i < num_items; ++i) + { + RandomBits(h_keys_in[i]); + } + + // Compute the reference solution: largest K of each segment, sorted descending + std::vector reference(static_cast(num_segments) * k); + for (int s = 0; s < num_segments; ++s) + { + std::vector seg(h_keys_in.begin() + s * segment_size, h_keys_in.begin() + (s + 1) * segment_size); + std::partial_sort(seg.begin(), seg.begin() + k, seg.end(), std::greater{}); + std::copy(seg.begin(), seg.begin() + k, reference.begin() + static_cast(s) * k); + } + + // Allocate device arrays + thrust::device_vector d_keys_in_buffer(h_keys_in); + thrust::device_vector d_keys_out_buffer(static_cast(num_segments) * k, thrust::no_init); + + // Build per-segment iterators: d_keys_in[s] points to the start of segment s. + auto d_keys_in = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(d_keys_in_buffer.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(d_keys_out_buffer.data())), k); + + // Specify that we do not require a specific output order and do not require deterministic results + auto requirements = + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted); + cudaStream_t stream = nullptr; + CubDebugExit(cudaStreamCreate(&stream)); + auto env = cuda::std::execution::env{cuda::stream_ref{stream}, requirements}; + + // Annotate the arguments: segment size and K are compile-time constants; the number of segments and the upper bound + // on the total number of items are runtime values. + auto segment_sizes = cuda::args::constant{}; + auto k_arg = cuda::args::constant{}; + auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; + auto total_num_args = cuda::args::immediate{static_cast(num_items)}; + + // Query temporary storage requirements + size_t temp_storage_bytes = 0; + CubDebugExit(DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs_arg, total_num_args, env)); + + // Allocate temporary storage + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + void* d_temp_storage = thrust::raw_pointer_cast(temp_storage.data()); + + // Run the segmented top-k algorithm + CubDebugExit(DeviceBatchedTopK::MaxKeys( + d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs_arg, total_num_args, env)); + + // Check for correctness: the per-segment output is unordered, so sort each output segment descending before + // comparing against the (descending) reference. + thrust::host_vector h_keys_out(d_keys_out_buffer); + int compare = 0; + for (int s = 0; s < num_segments && compare == 0; ++s) + { + std::sort(h_keys_out.begin() + s * k, h_keys_out.begin() + (s + 1) * k, std::greater{}); + compare = CompareResults( + reference.data() + static_cast(s) * k, + thrust::raw_pointer_cast(h_keys_out.data()) + static_cast(s) * k, + k, + g_verbose); + } + AssertEquals(0, compare); + + CubDebugExit(cudaStreamDestroy(stream)); + std::cout << "\n\n"; + return 0; +} diff --git a/cub/examples/device/example_device_batched_topk_pairs.cu b/cub/examples/device/example_device_batched_topk_pairs.cu new file mode 100644 index 000000000000..0db288e68261 --- /dev/null +++ b/cub/examples/device/example_device_batched_topk_pairs.cu @@ -0,0 +1,158 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! Simple example of cub::DeviceBatchedTopK::MaxPairs(). +//! For each (small) segment, find the K largest float keys and gather their associated values. This example +//! demonstrates the argument annotation framework with *runtime* values that carry a *compile-time* upper bound: +//! the segment size and K are passed as `cuda::args::immediate{value, cuda::args::bounds<1, MAX>()}`. +//! +//! To compile using the command line: +//! nvcc -arch=sm_XX example_device_batched_topk_pairs.cu -I../.. -lcudart -O3 + +// Ensure printing of CUDA runtime errors to console +#define CUB_STDERR + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "../../test/test_util.h" + +using namespace cub; + +bool g_verbose = false; + +// Compile-time upper bounds on the (small) segment size and K. The actual values are provided at runtime but are +// guaranteed not to exceed these bounds, which lets cub::DeviceBatchedTopK specialize for a single block per segment. +using offset_t = cuda::std::int64_t; +static constexpr offset_t max_seg_size = 1024; +static constexpr offset_t max_k = 32; + +int main(int argc, char** argv) +{ + offset_t segment_size = 256; + offset_t k = 8; + int num_segments = 2000; + + CommandLineArgs args(argc, argv); + g_verbose = args.CheckCmdLineFlag("v"); + args.GetCmdLineArgument("segment-size", segment_size); + args.GetCmdLineArgument("k", k); + args.GetCmdLineArgument("num-segments", num_segments); + + if (args.CheckCmdLineFlag("help")) + { + std::cout << "Usage: " << argv[0] + << " [--segment-size=] [--k=] [--num-segments=] [--device=] [--v]\n"; + exit(0); + } + CubDebugExit(args.DeviceInit()); + + segment_size = std::min(segment_size, max_seg_size); + k = std::min(k, segment_size); + const offset_t num_items = static_cast(num_segments) * segment_size; + + std::cout << "cub::DeviceBatchedTopK::MaxPairs() finds the " << k << " largest items in each of " << num_segments + << " segments of " << segment_size << " items\n"; + + // Initialize host input keys; values are the global item indices. + thrust::host_vector h_keys_in(num_items); + for (offset_t i = 0; i < num_items; ++i) + { + RandomBits(h_keys_in[i]); + } + + thrust::device_vector d_keys_in_buffer(h_keys_in); + thrust::device_vector d_keys_out_buffer(static_cast(num_segments) * k, thrust::no_init); + thrust::device_vector d_values_out_buffer(static_cast(num_segments) * k, thrust::no_init); + + // Per-segment iterators. Input values are the global indices [0, num_items) via a counting iterator. + auto d_keys_in = cuda::make_strided_iterator( + cuda::make_counting_iterator(thrust::raw_pointer_cast(d_keys_in_buffer.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(d_keys_out_buffer.data())), k); + auto d_values_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(offset_t{0})), segment_size); + auto d_values_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(d_values_out_buffer.data())), k); + + auto requirements = + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted); + cudaStream_t stream = nullptr; + CubDebugExit(cudaStreamCreate(&stream)); + auto env = cuda::std::execution::env{cuda::stream_ref{stream}, requirements}; + + // Annotate the arguments: runtime segment size and K, each with a compile-time upper bound. + auto segment_sizes = cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; + auto total_num_args = cuda::args::immediate{num_items}; + + size_t temp_storage_bytes = 0; + CubDebugExit(DeviceBatchedTopK::MaxPairs( + nullptr, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs_arg, + total_num_args, + env)); + + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + CubDebugExit(DeviceBatchedTopK::MaxPairs( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs_arg, + total_num_args, + env)); + + // Validate: (1) each output value indexes back to the matching output key, and (2) the selected keys are the K + // largest of their segment. + thrust::host_vector h_keys_out(d_keys_out_buffer); + thrust::host_vector h_values_out(d_values_out_buffer); + bool ok = true; + for (offset_t s = 0; s < num_segments && ok; ++s) + { + // (1) value/key association + for (offset_t j = 0; j < k; ++j) + { + const offset_t global_idx = h_values_out[s * k + j]; + ok = ok && global_idx >= 0 && global_idx < num_items && h_keys_in[global_idx] == h_keys_out[s * k + j]; + } + // (2) key correctness against a host reference (top-k largest, order-independent) + std::vector seg(h_keys_in.begin() + s * segment_size, h_keys_in.begin() + (s + 1) * segment_size); + std::partial_sort(seg.begin(), seg.begin() + k, seg.end(), std::greater{}); + std::vector res(h_keys_out.begin() + s * k, h_keys_out.begin() + (s + 1) * k); + std::sort(res.begin(), res.end(), std::greater{}); + ok = ok && std::equal(seg.begin(), seg.begin() + k, res.begin()); + } + AssertEquals(true, ok); + + CubDebugExit(cudaStreamDestroy(stream)); + std::cout << "\n\n"; + return 0; +} diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu new file mode 100644 index 000000000000..eeade32d7839 --- /dev/null +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -0,0 +1,259 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "insert_nested_NVTX_range_guard.h" + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +// Two segments of eight keys each. The per-segment top-3 max are {8,7,6} and {9,8,7}; the top-3 min are {-3,1,2} and +// {0,1,2}. +static auto make_two_segments() +{ + return thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; +} + +C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_topk][device]") +{ + // example-begin batched-topk-max-keys + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = make_two_segments(); + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + + // Per-segment iterators: d_keys_in[s] yields an iterator to the start of segment s. + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + // Argument annotations: a small, compile-time segment size and k, plus the runtime segment count and item-count + // bound. + auto segment_sizes = cuda::args::constant{}; + auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto total_items = cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}; + + // Top-k output is unordered and may be non-deterministic; this must be acknowledged via the environment. + auto env = cuda::std::execution::env{ + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + + // Query temporary storage requirements + size_t temp_storage_bytes = 0; + cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, total_items, env); + + // Allocate temporary storage and run + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + total_items, + env); + // example-end batched-topk-max-keys + + // The per-segment output is unordered; sort each segment (descending) before comparing. + thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); +} + +C2H_TEST("cub::DeviceBatchedTopK::MinKeys temp-storage API example", "[batched_topk][device]") +{ + // example-begin batched-topk-min-keys + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = make_two_segments(); + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + auto segment_sizes = cuda::args::constant{}; + auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto total_items = cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}; + auto env = cuda::std::execution::env{ + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + + size_t temp_storage_bytes = 0; + cub::DeviceBatchedTopK::MinKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, total_items, env); + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + cub::DeviceBatchedTopK::MinKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + total_items, + env); + // example-end batched-topk-min-keys + + thrust::sort(keys_out.begin(), keys_out.begin() + k); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); + REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); +} + +C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_topk][device]") +{ + // example-begin batched-topk-max-pairs + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = make_two_segments(); + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + auto values_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + // Input values are the per-segment item indices [0, segment_size). + auto d_values_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(0)), segment_size); + auto d_values_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); + + auto segment_sizes = cuda::args::constant{}; + auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto total_items = cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}; + auto env = cuda::std::execution::env{ + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + + size_t temp_storage_bytes = 0; + cub::DeviceBatchedTopK::MaxPairs( + nullptr, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs, + total_items, + env); + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + cub::DeviceBatchedTopK::MaxPairs( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs, + total_items, + env); + // example-end batched-topk-max-pairs + + // Each input value is the global item index, so each returned value indexes back to its returned key. + thrust::host_vector h_keys_in(keys_in); + thrust::host_vector h_keys_out(keys_out); + thrust::host_vector h_values_out(values_out); + for (int s = 0; s < num_segments; ++s) + { + for (int j = 0; j < k; ++j) + { + REQUIRE(h_keys_in[h_values_out[s * k + j]] == h_keys_out[s * k + j]); + } + } +} + +C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_topk][device]") +{ + // example-begin batched-topk-min-pairs + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = make_two_segments(); + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + auto values_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + auto d_values_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(0)), segment_size); + auto d_values_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); + + auto segment_sizes = cuda::args::constant{}; + auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto total_items = cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}; + auto env = cuda::std::execution::env{ + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + + size_t temp_storage_bytes = 0; + cub::DeviceBatchedTopK::MinPairs( + nullptr, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs, + total_items, + env); + thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); + cub::DeviceBatchedTopK::MinPairs( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs, + total_items, + env); + // example-end batched-topk-min-pairs + + thrust::host_vector h_keys_in(keys_in); + thrust::host_vector h_keys_out(keys_out); + thrust::host_vector h_values_out(values_out); + for (int s = 0; s < num_segments; ++s) + { + for (int j = 0; j < k; ++j) + { + REQUIRE(h_keys_in[h_values_out[s * k + j]] == h_keys_out[s * k + j]); + } + } +} diff --git a/cub/test/catch2_test_device_batched_topk_env_api.cu b/cub/test/catch2_test_device_batched_topk_env_api.cu new file mode 100644 index 000000000000..25d7ab1bc2cc --- /dev/null +++ b/cub/test/catch2_test_device_batched_topk_env_api.cu @@ -0,0 +1,211 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#include "insert_nested_NVTX_range_guard.h" + +#include + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +static auto make_two_segments() +{ + return thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; +} + +C2H_TEST("cub::DeviceBatchedTopK::MaxKeys env-alloc example", "[batched_topk][device][env]") +{ + // example-begin batched-topk-max-keys-env + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = make_two_segments(); + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + cuda::stream stream{cuda::devices[0]}; + auto env = cuda::std::execution::env{ + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted), + cuda::stream_ref{stream}}; + + // The env-based overload allocates and frees the temporary storage internally. + auto error = cub::DeviceBatchedTopK::MaxKeys( + d_keys_in, + d_keys_out, + cuda::args::constant{}, + cuda::args::constant{}, + cuda::args::immediate{cuda::std::int64_t{num_segments}}, + cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}, + env); + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MaxKeys failed with status: " << error << '\n'; + } + // example-end batched-topk-max-keys-env + + stream.sync(); + REQUIRE(error == cudaSuccess); + thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); +} + +C2H_TEST("cub::DeviceBatchedTopK::MinKeys env-alloc example", "[batched_topk][device][env]") +{ + // example-begin batched-topk-min-keys-env + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = make_two_segments(); + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + cuda::stream stream{cuda::devices[0]}; + auto env = cuda::std::execution::env{ + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted), + cuda::stream_ref{stream}}; + + auto error = cub::DeviceBatchedTopK::MinKeys( + d_keys_in, + d_keys_out, + cuda::args::constant{}, + cuda::args::constant{}, + cuda::args::immediate{cuda::std::int64_t{num_segments}}, + cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}, + env); + // example-end batched-topk-min-keys-env + + stream.sync(); + REQUIRE(error == cudaSuccess); + thrust::sort(keys_out.begin(), keys_out.begin() + k); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); + REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); +} + +C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][device][env]") +{ + // example-begin batched-topk-max-pairs-env + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = make_two_segments(); + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + auto values_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + auto d_values_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(0)), segment_size); + auto d_values_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); + + cuda::stream stream{cuda::devices[0]}; + auto env = cuda::std::execution::env{ + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted), + cuda::stream_ref{stream}}; + + auto error = cub::DeviceBatchedTopK::MaxPairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::constant{}, + cuda::args::constant{}, + cuda::args::immediate{cuda::std::int64_t{num_segments}}, + cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}, + env); + // example-end batched-topk-max-pairs-env + + stream.sync(); + REQUIRE(error == cudaSuccess); + thrust::host_vector h_keys_in(keys_in); + thrust::host_vector h_keys_out(keys_out); + thrust::host_vector h_values_out(values_out); + for (int s = 0; s < num_segments; ++s) + { + for (int j = 0; j < k; ++j) + { + REQUIRE(h_keys_in[h_values_out[s * k + j]] == h_keys_out[s * k + j]); + } + } +} + +C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][device][env]") +{ + // example-begin batched-topk-min-pairs-env + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = make_two_segments(); + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + auto values_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + auto d_values_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(0)), segment_size); + auto d_values_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); + + cuda::stream stream{cuda::devices[0]}; + auto env = cuda::std::execution::env{ + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted), + cuda::stream_ref{stream}}; + + auto error = cub::DeviceBatchedTopK::MinPairs( + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + cuda::args::constant{}, + cuda::args::constant{}, + cuda::args::immediate{cuda::std::int64_t{num_segments}}, + cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}, + env); + // example-end batched-topk-min-pairs-env + + stream.sync(); + REQUIRE(error == cudaSuccess); + thrust::host_vector h_keys_in(keys_in); + thrust::host_vector h_keys_out(keys_out); + thrust::host_vector h_values_out(values_out); + for (int s = 0; s < num_segments; ++s) + { + for (int j = 0; j < k; ++j) + { + REQUIRE(h_keys_in[h_values_out[s * k + j]] == h_keys_out[s * k + j]); + } + } +} diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 20f2f7937831..c1c50019f9b0 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -3,13 +3,16 @@ #include "insert_nested_NVTX_range_guard.h" -#include +#include #include #include #include #include +#include +#include +#include #include #include @@ -27,7 +30,7 @@ struct is_minus_zero } }; -template (nullptr); - return cub::detail::batched_topk::dispatch( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - values_it, - values_it, - segment_sizes, - k, - cuda::args::constant{}, - num_segments, - total_num_items_guarantee, - stream); + auto env = cuda::std::execution::env{ + cuda::stream_ref{stream}, + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + if constexpr (SelectDirection == cub::detail::topk::select::max) + { + return cub::DeviceBatchedTopK::MaxKeys( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k, + num_segments, + total_num_items_guarantee, + env); + } + else + { + return cub::DeviceBatchedTopK::MinKeys( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + segment_sizes, + k, + num_segments, + total_num_items_guarantee, + env); + } } // %PARAM% TEST_LAUNCH lid 0:1:2 DECLARE_TMPL_LAUNCH_WRAPPER( - dispatch_batched_topk_keys, batched_topk_keys, cub::detail::topk::select Direction, Direction); + dispatch_batched_topk_keys, batched_topk_keys, cub::detail::topk::select SelectDirection, SelectDirection); // Total segment size using max_segment_size_list = c2h::enum_type_list; diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 9e01d126896c..db1607f5a3a8 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -3,13 +3,16 @@ #include "insert_nested_NVTX_range_guard.h" -#include +#include #include #include #include #include +#include +#include +#include #include #include "catch2_test_device_topk_common.cuh" @@ -47,46 +50,66 @@ struct flag_intra_segment_duplicates template flag_intra_segment_duplicates(ItemItT, SegIdItT) -> flag_intra_segment_duplicates; -template CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( void* d_temp_storage, size_t& temp_storage_bytes, - KeyInputItItT d_key_segments_it, - KeyOutputItItT d_key_segments_out_it, - ValueInputItItT d_value_segments_it, - ValueOutputItItT d_value_segments_out_it, - SegmentSizeParameterT segment_sizes, - KParameterT k, + KeyInputItItT d_keys_in, + KeyOutputItItT d_keys_out, + ValueInputItItT d_values_in, + ValueOutputItItT d_values_out, + SegmentSizeParamT segment_sizes, + KParamT k, NumSegmentsParameterT num_segments, TotalNumItemsGuaranteeT total_num_items_guarantee, cudaStream_t stream = nullptr) { - return cub::detail::batched_topk::dispatch( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - d_value_segments_it, - d_value_segments_out_it, - segment_sizes, - k, - cuda::args::constant{}, - num_segments, - total_num_items_guarantee, - stream); + auto env = cuda::std::execution::env{ + cuda::stream_ref{stream}, + cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + if constexpr (SelectDirection == cub::detail::topk::select::max) + { + return cub::DeviceBatchedTopK::MaxPairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + num_segments, + total_num_items_guarantee, + env); + } + else + { + return cub::DeviceBatchedTopK::MinPairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k, + num_segments, + total_num_items_guarantee, + env); + } } // %PARAM% TEST_LAUNCH lid 0:1:2 DECLARE_TMPL_LAUNCH_WRAPPER( - dispatch_batched_topk_pairs, batched_topk_pairs, cub::detail::topk::select Direction, Direction); + dispatch_batched_topk_pairs, batched_topk_pairs, cub::detail::topk::select SelectDirection, SelectDirection); // Total segment size using max_segment_size_list = c2h::enum_type_list; From 8c8e07a08acf11f9da89166d59f3be5edabea3fb Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Sun, 14 Jun 2026 13:31:02 -0700 Subject: [PATCH 02/18] Document DeviceBatchedTopK determinism, tie-break, and output ordering Update cub::DeviceBatchedTopK to document the committed default contract (determinism::run_to_run, tie_break::prefer_smaller_index, output_ordering::stable_sorted) and to note the currently-supported, explicitly-requested non-deterministic, unsorted configuration. Links the top-k execution-requirements page (:ref:`cub-topk-requirements`), which is added separately on the docs/topk-determinism branch; this is a forward reference until that docs PR merges and this branch is rebased. --- cub/cub/device/device_batched_topk.cuh | 32 +++++++++++++++++++++----- 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 8d798272d609..766a4220a893 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -195,15 +195,35 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! enough that such a block fits within the shared-memory limit. Both uniform (fixed) and variable segment sizes are //! supported as long as this maximum is honored. //! - **Uniform number of segments.** ``num_segments`` must be a single (host-resolved) value. -//! - **Non-deterministic, unsorted output.** Like :cpp:struct:`cub::DeviceTopK`, the caller must acknowledge this by -//! passing ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, +//! - **Explicit opt-out required for the output guarantees.** The deterministic, stable-sorted default contract +//! described in *Determinism, tie-breaking, and output ordering* below (and in :ref:`cub-topk-requirements`) is not +//! yet implemented. Like :cpp:struct:`cub::DeviceTopK`, the caller must currently request non-deterministic, +//! unsorted output explicitly by passing ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, //! cuda::execution::output_ordering::unsorted)`` in the environment. //! -//! Determinism -//! ++++++++++++++++++++++++++ +//! Determinism, tie-breaking, and output ordering +//! +++++++++++++++++++++++++++++++++++++++++++++++ +//! +//! Like :cpp:struct:`cub::DeviceTopK`, the result of ``DeviceBatchedTopK`` is governed by two orthogonal execution +//! requirements: *which* items are selected per segment (``cuda::execution::determinism``, optionally refined by +//! ``cuda::execution::tie_break``) and the order in which they are written (``cuda::execution::output_ordering``). +//! When the caller does not opt out, the committed default is the most reproducible behavior: deterministic results +//! (``cuda::execution::determinism::run_to_run``), ties resolved toward the smaller (lower) source index +//! (``cuda::execution::tie_break::prefer_smaller_index``), and stable-sorted output +//! (``cuda::execution::output_ordering::stable_sorted``). Callers opt *out* of these guarantees to obtain faster +//! implementations. +//! +//! See :ref:`cub-topk-requirements` for the full requirement model, worked examples, and guidance on choosing +//! requirements. +//! +//! .. note:: //! -//! DeviceBatchedTopK currently only supports unordered output, which may be non-deterministic for certain inputs: -//! if multiple items tie at the K-th position, the subset of tied elements returned may vary between runs. +//! **Current support.** This release only implements the fully opted-out configuration, which must be requested +//! explicitly: ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, +//! cuda::execution::output_ordering::unsorted)``. Any other combination (including an empty, no-requirement +//! environment) is rejected at compile time. In this configuration the per-segment output is unordered and may be +//! non-deterministic: if multiple items tie at the K-th position, the subset of tied elements returned is not +//! uniquely defined and may vary between runs. //! //! Usage Considerations //! ++++++++++++++++++++++++++ From 536e0e265531902ccf8a3f18408c65d15a7d36bd Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Mon, 15 Jun 2026 02:51:14 -0700 Subject: [PATCH 03/18] extends top-k requirements page for DeviceBatchedTopK --- docs/cub/api_docs/device_topk_requirements.rst | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/docs/cub/api_docs/device_topk_requirements.rst b/docs/cub/api_docs/device_topk_requirements.rst index 11d95d314550..9b02280fea4b 100644 --- a/docs/cub/api_docs/device_topk_requirements.rst +++ b/docs/cub/api_docs/device_topk_requirements.rst @@ -5,9 +5,11 @@ Top-K: Determinism, Tie-Breaking, and Output Ordering ====================================================== -This page describes how to control the result of the CUB top-k family of algorithms (currently -:cpp:struct:`cub::DeviceTopK`) through the execution environment. The same requirement model -applies to every ``MaxKeys`` / ``MinKeys`` / ``MaxPairs`` / ``MinPairs`` entry point. +This page describes how to control the result of the CUB top-k family of algorithms +(:cpp:struct:`cub::DeviceTopK` and :cpp:struct:`cub::DeviceBatchedTopK`) through the execution +environment. For :cpp:struct:`cub::DeviceBatchedTopK`, these requirements apply independently within +each segment. The same requirement model applies to every ``MaxKeys`` / ``MinKeys`` / ``MaxPairs`` / +``MinPairs`` entry point. Two orthogonal concerns ----------------------- From bad1e4c226865b53cc0c56a69595804ac6eeb7a7 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 16 Jun 2026 07:21:25 -0700 Subject: [PATCH 04/18] improves docs and examples --- cub/cub/device/device_batched_topk.cuh | 64 ++---- .../example_device_batched_topk_keys.cu | 14 +- .../example_device_batched_topk_pairs.cu | 5 +- .../catch2_test_device_batched_topk_api.cu | 186 +++++++++++------- ...catch2_test_device_batched_topk_env_api.cu | 63 ++++-- .../catch2_test_device_segmented_topk_keys.cu | 21 +- ...catch2_test_device_segmented_topk_pairs.cu | 18 +- 7 files changed, 194 insertions(+), 177 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 766a4220a893..ef3df8d9d542 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -33,6 +33,8 @@ #include #include #include +#include +#include CUB_NAMESPACE_BEGIN @@ -56,7 +58,6 @@ template CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( void* d_temp_storage, @@ -68,7 +69,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env) { // --------------------------------------------------------------------------- @@ -128,7 +128,12 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( "host; pass num_segments as a single-value annotation (e.g. cuda::args::constant or " "cuda::args::immediate), not a per-segment sequence."); - auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); + const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); + + // The total-number-of-items guarantee is intentionally not part of the initial public API surface. The dispatch + // only uses its element type to size internal large-segment offsets (the value itself is unused), so we pass a + // conservative 64-bit upper bound here. + const auto total_num_items = ::cuda::args::immediate{::cuda::std::numeric_limits<::cuda::std::int64_t>::max()}; return batched_topk::dispatch( d_temp_storage, @@ -156,21 +161,20 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! ++++++++++++++++++++++++++ //! //! Given a batch of segments, ``DeviceBatchedTopK`` finds, independently for each segment, the K largest (or -//! smallest) items. Within each output segment the K results are returned in an **unspecified order**. The algorithm -//! reuses the block-level top-k primitive used by :cpp:struct:`cub::DeviceTopK`. +//! smallest) items. //! //! Argument annotation framework //! +++++++++++++++++++++++++++++++++++++++++++++ //! -//! The per-segment parameters (segment sizes, ``k``, the number of segments, and a guarantee on the total number of -//! input items) are supplied as **annotated arguments** from ``cuda::args``. This lets callers communicate +//! The per-segment parameters (segment sizes, ``k``, and the number of segments) are supplied as **annotated +//! arguments** from ``cuda::args``. This lets callers communicate //! compile-time guarantees (e.g. a static upper bound on the segment size) that the algorithm uses to specialize the //! kernel: //! -//! - ``cuda::args::constant`` — a value known at compile time. -//! - ``cuda::args::immediate{value, cuda::args::bounds()}`` — a runtime value with a +//! - ``cuda::args::constant`` for a value known at compile time. +//! - ``cuda::args::immediate{value, cuda::args::bounds()}`` for a runtime value with a //! compile-time ``[lo, hi]`` range. -//! - ``cuda::args::deferred_sequence{iterator, cuda::args::bounds()}`` — a per-segment +//! - ``cuda::args::deferred_sequence{iterator, cuda::args::bounds()}`` for a per-segment //! sequence (e.g. variable segment sizes) with a compile-time element range. //! //! Choosing argument bounds @@ -277,8 +281,8 @@ struct DeviceBatchedTopK //! //! @param[in] segment_sizes //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, - //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum; - //! prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the //! *Choosing argument bounds* section). //! //! @param[in] k @@ -287,10 +291,6 @@ struct DeviceBatchedTopK //! @param[in] num_segments //! Annotated argument providing the (uniform) number of segments //! - //! @param[in] total_num_items - //! Annotated argument providing an upper-bound guarantee on the total number of input items across all segments, - //! used to size internal offset types - //! //! @param[in] env //! @rst //! **[optional]** Execution environment. Must require `determinism::not_guaranteed` and @@ -301,9 +301,8 @@ struct DeviceBatchedTopK typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, typename EnvT = ::cuda::std::execution::env<>> - CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys( + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys( void* d_temp_storage, size_t& temp_storage_bytes, KeyInputIteratorItT d_keys_in, @@ -311,7 +310,6 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MaxKeys"); @@ -325,7 +323,6 @@ struct DeviceBatchedTopK segment_sizes, k, num_segments, - total_num_items, ::cuda::std::move(env)); } @@ -354,7 +351,6 @@ struct DeviceBatchedTopK typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, typename EnvT = ::cuda::std::execution::env<>> [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys( KeyInputIteratorItT d_keys_in, @@ -362,7 +358,6 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxKeys"); @@ -378,7 +373,6 @@ struct DeviceBatchedTopK segment_sizes, k, num_segments, - total_num_items, env); }); } @@ -409,9 +403,8 @@ struct DeviceBatchedTopK typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, typename EnvT = ::cuda::std::execution::env<>> - CUB_RUNTIME_FUNCTION static cudaError_t MinKeys( + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinKeys( void* d_temp_storage, size_t& temp_storage_bytes, KeyInputIteratorItT d_keys_in, @@ -419,7 +412,6 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MinKeys"); @@ -433,7 +425,6 @@ struct DeviceBatchedTopK segment_sizes, k, num_segments, - total_num_items, ::cuda::std::move(env)); } @@ -460,7 +451,6 @@ struct DeviceBatchedTopK typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, typename EnvT = ::cuda::std::execution::env<>> [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinKeys( KeyInputIteratorItT d_keys_in, @@ -468,7 +458,6 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinKeys"); @@ -484,7 +473,6 @@ struct DeviceBatchedTopK segment_sizes, k, num_segments, - total_num_items, env); }); } @@ -534,9 +522,8 @@ struct DeviceBatchedTopK typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, typename EnvT = ::cuda::std::execution::env<>> - CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs( + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs( void* d_temp_storage, size_t& temp_storage_bytes, KeyInputIteratorItT d_keys_in, @@ -546,7 +533,6 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MaxPairs"); @@ -560,7 +546,6 @@ struct DeviceBatchedTopK segment_sizes, k, num_segments, - total_num_items, ::cuda::std::move(env)); } @@ -585,7 +570,6 @@ struct DeviceBatchedTopK typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, typename EnvT = ::cuda::std::execution::env<>> [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs( KeyInputIteratorItT d_keys_in, @@ -595,7 +579,6 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxPairs"); @@ -611,7 +594,6 @@ struct DeviceBatchedTopK segment_sizes, k, num_segments, - total_num_items, env); }); } @@ -645,9 +627,8 @@ struct DeviceBatchedTopK typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, typename EnvT = ::cuda::std::execution::env<>> - CUB_RUNTIME_FUNCTION static cudaError_t MinPairs( + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinPairs( void* d_temp_storage, size_t& temp_storage_bytes, KeyInputIteratorItT d_keys_in, @@ -657,7 +638,6 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MinPairs"); @@ -671,7 +651,6 @@ struct DeviceBatchedTopK segment_sizes, k, num_segments, - total_num_items, ::cuda::std::move(env)); } @@ -696,7 +675,6 @@ struct DeviceBatchedTopK typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, - typename TotalNumItemsGuaranteeT, typename EnvT = ::cuda::std::execution::env<>> [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinPairs( KeyInputIteratorItT d_keys_in, @@ -706,7 +684,6 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items, EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinPairs"); @@ -722,7 +699,6 @@ struct DeviceBatchedTopK segment_sizes, k, num_segments, - total_num_items, env); }); } diff --git a/cub/examples/device/example_device_batched_topk_keys.cu b/cub/examples/device/example_device_batched_topk_keys.cu index 605552300dc3..6bbf10054bbb 100644 --- a/cub/examples/device/example_device_batched_topk_keys.cu +++ b/cub/examples/device/example_device_batched_topk_keys.cu @@ -96,17 +96,15 @@ int main(int argc, char** argv) CubDebugExit(cudaStreamCreate(&stream)); auto env = cuda::std::execution::env{cuda::stream_ref{stream}, requirements}; - // Annotate the arguments: segment size and K are compile-time constants; the number of segments and the upper bound - // on the total number of items are runtime values. - auto segment_sizes = cuda::args::constant{}; - auto k_arg = cuda::args::constant{}; - auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; - auto total_num_args = cuda::args::immediate{static_cast(num_items)}; + // Annotate the arguments: segment size and K are compile-time constants; the number of segments is a runtime value. + auto segment_sizes = cuda::args::constant{}; + auto k_arg = cuda::args::constant{}; + auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; // Query temporary storage requirements size_t temp_storage_bytes = 0; CubDebugExit(DeviceBatchedTopK::MaxKeys( - nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs_arg, total_num_args, env)); + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs_arg, env)); // Allocate temporary storage thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); @@ -114,7 +112,7 @@ int main(int argc, char** argv) // Run the segmented top-k algorithm CubDebugExit(DeviceBatchedTopK::MaxKeys( - d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs_arg, total_num_args, env)); + d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs_arg, env)); // Check for correctness: the per-segment output is unordered, so sort each output segment descending before // comparing against the (descending) reference. diff --git a/cub/examples/device/example_device_batched_topk_pairs.cu b/cub/examples/device/example_device_batched_topk_pairs.cu index 0db288e68261..aa7adbdd7a89 100644 --- a/cub/examples/device/example_device_batched_topk_pairs.cu +++ b/cub/examples/device/example_device_batched_topk_pairs.cu @@ -99,8 +99,7 @@ int main(int argc, char** argv) // Annotate the arguments: runtime segment size and K, each with a compile-time upper bound. auto segment_sizes = cuda::args::immediate{segment_size, cuda::args::bounds()}; auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; - auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; - auto total_num_args = cuda::args::immediate{num_items}; + auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; size_t temp_storage_bytes = 0; CubDebugExit(DeviceBatchedTopK::MaxPairs( @@ -113,7 +112,6 @@ int main(int argc, char** argv) segment_sizes, k_arg, num_segs_arg, - total_num_args, env)); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); @@ -127,7 +125,6 @@ int main(int argc, char** argv) segment_sizes, k_arg, num_segs_arg, - total_num_args, env)); // Validate: (1) each output value indexes back to the matching output key, and (2) the selected keys are the K diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu index eeade32d7839..3365715236c5 100644 --- a/cub/test/catch2_test_device_batched_topk_api.cu +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -18,14 +18,9 @@ #include #include -#include +#include -// Two segments of eight keys each. The per-segment top-3 max are {8,7,6} and {9,8,7}; the top-3 min are {-3,1,2} and -// {0,1,2}. -static auto make_two_segments() -{ - return thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; -} +#include C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_topk][device]") { @@ -34,7 +29,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_t constexpr int segment_size = 8; constexpr int k = 3; - auto keys_in = make_two_segments(); + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); // Per-segment iterators: d_keys_in[s] yields an iterator to the start of segment s. @@ -48,7 +43,6 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_t auto segment_sizes = cuda::args::constant{}; auto k_arg = cuda::args::constant{}; auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto total_items = cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}; // Top-k output is unordered and may be non-deterministic; this must be acknowledged via the environment. auto env = cuda::std::execution::env{ @@ -56,23 +50,30 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_t // Query temporary storage requirements size_t temp_storage_bytes = 0; - cub::DeviceBatchedTopK::MaxKeys( - nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, total_items, env); + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); // Allocate temporary storage and run thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - cub::DeviceBatchedTopK::MaxKeys( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - segment_sizes, - k_arg, - num_segs, - total_items, - env); + if (error == cudaSuccess) + { + error = cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); + } + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MaxKeys failed with status: " << error << '\n'; + } // example-end batched-topk-max-keys + REQUIRE(error == cudaSuccess); // The per-segment output is unordered; sort each segment (descending) before comparing. thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); @@ -86,7 +87,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys temp-storage API example", "[batched_t constexpr int segment_size = 8; constexpr int k = 3; - auto keys_in = make_two_segments(); + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); auto d_keys_in = @@ -97,26 +98,32 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys temp-storage API example", "[batched_t auto segment_sizes = cuda::args::constant{}; auto k_arg = cuda::args::constant{}; auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto total_items = cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}; auto env = cuda::std::execution::env{ cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; size_t temp_storage_bytes = 0; - cub::DeviceBatchedTopK::MinKeys( - nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, total_items, env); + auto error = cub::DeviceBatchedTopK::MinKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - cub::DeviceBatchedTopK::MinKeys( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - segment_sizes, - k_arg, - num_segs, - total_items, - env); + if (error == cudaSuccess) + { + error = cub::DeviceBatchedTopK::MinKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); + } + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MinKeys failed with status: " << error << '\n'; + } // example-end batched-topk-min-keys + REQUIRE(error == cudaSuccess); thrust::sort(keys_out.begin(), keys_out.begin() + k); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); @@ -129,7 +136,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ constexpr int segment_size = 8; constexpr int k = 3; - auto keys_in = make_two_segments(); + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); auto values_out = thrust::device_vector(num_segments * k, thrust::no_init); @@ -138,20 +145,18 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); // Input values are the per-segment item indices [0, segment_size). - auto d_values_in = - cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(0)), segment_size); + auto d_values_in = cuda::make_constant_iterator(cuda::make_counting_iterator(0)); auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); auto segment_sizes = cuda::args::constant{}; auto k_arg = cuda::args::constant{}; auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto total_items = cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}; auto env = cuda::std::execution::env{ cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; size_t temp_storage_bytes = 0; - cub::DeviceBatchedTopK::MaxPairs( + auto error = cub::DeviceBatchedTopK::MaxPairs( nullptr, temp_storage_bytes, d_keys_in, @@ -161,24 +166,32 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ segment_sizes, k_arg, num_segs, - total_items, env); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - cub::DeviceBatchedTopK::MaxPairs( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k_arg, - num_segs, - total_items, - env); + if (error == cudaSuccess) + { + error = cub::DeviceBatchedTopK::MaxPairs( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs, + env); + } + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MaxPairs failed with status: " << error << '\n'; + } // example-end batched-topk-max-pairs - // Each input value is the global item index, so each returned value indexes back to its returned key. + REQUIRE(error == cudaSuccess); + + // Each input value is the global item index. Bounds-checked: every returned value indexes back to the input + // element (within its own segment) whose key was selected. thrust::host_vector h_keys_in(keys_in); thrust::host_vector h_keys_out(keys_out); thrust::host_vector h_values_out(values_out); @@ -186,9 +199,18 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ { for (int j = 0; j < k; ++j) { - REQUIRE(h_keys_in[h_values_out[s * k + j]] == h_keys_out[s * k + j]); + const int idx = s * k + j; + const int v = h_values_out[idx]; + REQUIRE(v >= 0); + REQUIRE(v < segment_size); + REQUIRE(h_keys_in[s * segment_size + v] == h_keys_out[idx]); } } + + // The selected keys must be the per-segment top-k (output is unordered; sort each segment descending). + thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); } C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_topk][device]") @@ -198,7 +220,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ constexpr int segment_size = 8; constexpr int k = 3; - auto keys_in = make_two_segments(); + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); auto values_out = thrust::device_vector(num_segments * k, thrust::no_init); @@ -206,20 +228,18 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); - auto d_values_in = - cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(0)), segment_size); + auto d_values_in = cuda::make_constant_iterator(cuda::make_counting_iterator(0)); auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); auto segment_sizes = cuda::args::constant{}; auto k_arg = cuda::args::constant{}; auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto total_items = cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}; auto env = cuda::std::execution::env{ cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; size_t temp_storage_bytes = 0; - cub::DeviceBatchedTopK::MinPairs( + auto error = cub::DeviceBatchedTopK::MinPairs( nullptr, temp_storage_bytes, d_keys_in, @@ -229,23 +249,32 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ segment_sizes, k_arg, num_segs, - total_items, env); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - cub::DeviceBatchedTopK::MinPairs( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k_arg, - num_segs, - total_items, - env); + if (error == cudaSuccess) + { + error = cub::DeviceBatchedTopK::MinPairs( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs, + env); + } + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MinPairs failed with status: " << error << '\n'; + } // example-end batched-topk-min-pairs + REQUIRE(error == cudaSuccess); + + // Each input value is the global item index. Bounds-checked: every returned value indexes back to the input + // element (within its own segment) whose key was selected. thrust::host_vector h_keys_in(keys_in); thrust::host_vector h_keys_out(keys_out); thrust::host_vector h_values_out(values_out); @@ -253,7 +282,16 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ { for (int j = 0; j < k; ++j) { - REQUIRE(h_keys_in[h_values_out[s * k + j]] == h_keys_out[s * k + j]); + const int idx = s * k + j; + const int v = h_values_out[idx]; + REQUIRE(v >= 0); + REQUIRE(v < segment_size); + REQUIRE(h_keys_in[s * segment_size + v] == h_keys_out[idx]); } } + + // The selected keys must be the per-segment top-k (output is unordered; sort each segment ascending). + thrust::sort(keys_out.begin(), keys_out.begin() + k); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); + REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); } diff --git a/cub/test/catch2_test_device_batched_topk_env_api.cu b/cub/test/catch2_test_device_batched_topk_env_api.cu index 25d7ab1bc2cc..caa21e802cc8 100644 --- a/cub/test/catch2_test_device_batched_topk_env_api.cu +++ b/cub/test/catch2_test_device_batched_topk_env_api.cu @@ -24,11 +24,6 @@ #include -static auto make_two_segments() -{ - return thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; -} - C2H_TEST("cub::DeviceBatchedTopK::MaxKeys env-alloc example", "[batched_topk][device][env]") { // example-begin batched-topk-max-keys-env @@ -36,7 +31,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys env-alloc example", "[batched_topk][de constexpr int segment_size = 8; constexpr int k = 3; - auto keys_in = make_two_segments(); + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); auto d_keys_in = @@ -56,7 +51,6 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys env-alloc example", "[batched_topk][de cuda::args::constant{}, cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, - cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}, env); if (error != cudaSuccess) { @@ -78,7 +72,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys env-alloc example", "[batched_topk][de constexpr int segment_size = 8; constexpr int k = 3; - auto keys_in = make_two_segments(); + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); auto d_keys_in = @@ -97,8 +91,11 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys env-alloc example", "[batched_topk][de cuda::args::constant{}, cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, - cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}, env); + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MinKeys failed with status: " << error << '\n'; + } // example-end batched-topk-min-keys-env stream.sync(); @@ -115,7 +112,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][d constexpr int segment_size = 8; constexpr int k = 3; - auto keys_in = make_two_segments(); + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); auto values_out = thrust::device_vector(num_segments * k, thrust::no_init); @@ -123,8 +120,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][d cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); - auto d_values_in = - cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(0)), segment_size); + auto d_values_in = cuda::make_constant_iterator(cuda::make_counting_iterator(0)); auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); @@ -141,12 +137,18 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][d cuda::args::constant{}, cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, - cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}, env); + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MaxPairs failed with status: " << error << '\n'; + } // example-end batched-topk-max-pairs-env stream.sync(); REQUIRE(error == cudaSuccess); + + // Each input value is the global item index. Bounds-checked: every returned value indexes back to the input + // element (within its own segment) whose key was selected. thrust::host_vector h_keys_in(keys_in); thrust::host_vector h_keys_out(keys_out); thrust::host_vector h_values_out(values_out); @@ -154,9 +156,18 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][d { for (int j = 0; j < k; ++j) { - REQUIRE(h_keys_in[h_values_out[s * k + j]] == h_keys_out[s * k + j]); + const int idx = s * k + j; + const int v = h_values_out[idx]; + REQUIRE(v >= 0); + REQUIRE(v < segment_size); + REQUIRE(h_keys_in[s * segment_size + v] == h_keys_out[idx]); } } + + // The selected keys must be the per-segment top-k (output is unordered; sort each segment descending). + thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); } C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][device][env]") @@ -166,7 +177,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][d constexpr int segment_size = 8; constexpr int k = 3; - auto keys_in = make_two_segments(); + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); auto values_out = thrust::device_vector(num_segments * k, thrust::no_init); @@ -174,8 +185,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][d cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); - auto d_values_in = - cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(0)), segment_size); + auto d_values_in = cuda::make_constant_iterator(cuda::make_counting_iterator(0)); auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); @@ -192,12 +202,18 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][d cuda::args::constant{}, cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, - cuda::args::immediate{cuda::std::int64_t{num_segments * segment_size}}, env); + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MinPairs failed with status: " << error << '\n'; + } // example-end batched-topk-min-pairs-env stream.sync(); REQUIRE(error == cudaSuccess); + + // Each input value is the global item index. Bounds-checked: every returned value indexes back to the input + // element (within its own segment) whose key was selected. thrust::host_vector h_keys_in(keys_in); thrust::host_vector h_keys_out(keys_out); thrust::host_vector h_values_out(values_out); @@ -205,7 +221,16 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][d { for (int j = 0; j < k; ++j) { - REQUIRE(h_keys_in[h_values_out[s * k + j]] == h_keys_out[s * k + j]); + const int idx = s * k + j; + const int v = h_values_out[idx]; + REQUIRE(v >= 0); + REQUIRE(v < segment_size); + REQUIRE(h_keys_in[s * segment_size + v] == h_keys_out[idx]); } } + + // The selected keys must be the per-segment top-k (output is unordered; sort each segment ascending). + thrust::sort(keys_out.begin(), keys_out.begin() + k); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); + REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); } diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index c1c50019f9b0..542cdc4eba45 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -35,8 +35,7 @@ template + typename NumSegmentsParameterT> CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( void* d_temp_storage, size_t& temp_storage_bytes, @@ -45,7 +44,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( SegmentSizeParamT segment_sizes, KParamT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items_guarantee, cudaStream_t stream = nullptr) { auto env = cuda::std::execution::env{ @@ -61,7 +59,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( segment_sizes, k, num_segments, - total_num_items_guarantee, env); } else @@ -74,7 +71,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( segment_sizes, k, num_segments, - total_num_items_guarantee, env); } } @@ -186,8 +182,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small fixed-size segments", d_keys_out, cuda::args::immediate{segment_size, cuda::args::bounds()}, cuda::args::immediate{k, cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_segments * segment_size}); + cuda::args::immediate{num_segments}); // Prepare expected results fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); compact_sorted_keys_to_topk(expected_keys, segment_size, k); @@ -283,8 +278,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with small variable-size segment d_keys_out, cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}, cuda::args::immediate{k, cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_items}); + cuda::args::immediate{num_segments}); // Verify keys are returned correctly: sort each segment of the expected input, then compact the top-k segmented_sort_keys(expected_keys, num_segments, segment_offsets.cbegin(), segment_offsets.cbegin() + 1, direction); @@ -381,8 +375,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with fixed-size segments and per cuda::args::immediate{segment_size, cuda::args::bounds()}, cuda::args::deferred_sequence{ thrust::raw_pointer_cast(segment_k.data()), cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_segments * segment_size}); + cuda::args::immediate{num_segments}); // Prepare expected results: sort each fixed-size input segment, then compact each to its per-segment top-k. fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); @@ -480,8 +473,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys work with variable-size segments and cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}, cuda::args::deferred_sequence{ thrust::raw_pointer_cast(segment_k.data()), cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_items}); + cuda::args::immediate{num_segments}); // Verify keys are returned correctly: sort each segment of the expected input, then compact the per-segment top-k segmented_sort_keys(expected_keys, num_segments, segment_offsets.cbegin(), segment_offsets.cbegin() + 1, direction); @@ -523,8 +515,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys preserve -0.0f in output", d_keys_out_it, cuda::args::immediate{segment_size, cuda::args::bounds()}, cuda::args::immediate{k, cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_segments * segment_size}); + cuda::args::immediate{num_segments}); const int num_minus_zero = static_cast(thrust::count_if(d_keys_out.begin(), d_keys_out.end(), is_minus_zero{})); REQUIRE(num_minus_zero >= 1); diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index db1607f5a3a8..2143e479e3e3 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -57,8 +57,7 @@ template + typename NumSegmentsParameterT> CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( void* d_temp_storage, size_t& temp_storage_bytes, @@ -69,7 +68,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( SegmentSizeParamT segment_sizes, KParamT k, NumSegmentsParameterT num_segments, - TotalNumItemsGuaranteeT total_num_items_guarantee, cudaStream_t stream = nullptr) { auto env = cuda::std::execution::env{ @@ -87,7 +85,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( segment_sizes, k, num_segments, - total_num_items_guarantee, env); } else @@ -102,7 +99,6 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( segment_sizes, k, num_segments, - total_num_items_guarantee, env); } } @@ -292,8 +288,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small fixed-size segments" d_values_out, cuda::args::immediate{segment_size, cuda::args::bounds()}, cuda::args::immediate{k, cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_segments * segment_size}); + cuda::args::immediate{num_segments}); // Verification: // - We verify correct top-k selection through the keys @@ -412,8 +407,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with small variable-size segmen d_values_out, cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}, cuda::args::immediate{k, cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_items}); + cuda::args::immediate{num_segments}); // Verification: // - We verify correct top-k selection through the keys @@ -530,8 +524,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with fixed-size segments and pe cuda::args::immediate{segment_size, cuda::args::bounds()}, cuda::args::deferred_sequence{ thrust::raw_pointer_cast(segment_k.data()), cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_segments * segment_size}); + cuda::args::immediate{num_segments}); // Verification: // - We verify correct top-k selection through the keys @@ -650,8 +643,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Pairs work with variable-size segments and cuda::args::deferred_sequence{segment_size_it, cuda::args::bounds()}, cuda::args::deferred_sequence{ thrust::raw_pointer_cast(segment_k.data()), cuda::args::bounds()}, - cuda::args::immediate{num_segments}, - cuda::args::immediate{num_items}); + cuda::args::immediate{num_segments}); // Verification: // - We verify correct top-k selection through the keys From ca103b6da3137ad239088ae0fe18a33335e2a5b8 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 16 Jun 2026 22:44:26 -0700 Subject: [PATCH 05/18] better cuda::args docs --- cub/cub/device/device_batched_topk.cuh | 70 +++++++++++-------- .../example_device_batched_topk_pairs.cu | 6 +- .../catch2_test_device_batched_topk_api.cu | 22 +----- .../catch2_test_device_segmented_topk_keys.cu | 18 +---- 4 files changed, 47 insertions(+), 69 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index ef3df8d9d542..ee0038f7cdc5 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -166,16 +166,44 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! Argument annotation framework //! +++++++++++++++++++++++++++++++++++++++++++++ //! -//! The per-segment parameters (segment sizes, ``k``, and the number of segments) are supplied as **annotated -//! arguments** from ``cuda::args``. This lets callers communicate -//! compile-time guarantees (e.g. a static upper bound on the segment size) that the algorithm uses to specialize the -//! kernel: +//! The parameters ``segment_sizes``, ``k``, and ``num_segments`` are passed as **annotated arguments** from +//! ``cuda::args``. An annotation tells the algorithm everything you know about a parameter: where its value comes +//! from and how tightly it is bounded. The more you can tell the algorithm, and the more precisely (a +//! compile-time constant rather than a runtime value, a tight bound rather than a loose one), the more it can +//! specialize. For that reason, we encourage you to provide as much information as you have. //! -//! - ``cuda::args::constant`` for a value known at compile time. -//! - ``cuda::args::immediate{value, cuda::args::bounds()}`` for a runtime value with a -//! compile-time ``[lo, hi]`` range. -//! - ``cuda::args::deferred_sequence{iterator, cuda::args::bounds()}`` for a per-segment -//! sequence (e.g. variable segment sizes) with a compile-time element range. +//! **Where the value comes from.** The first three forms describe a single value shared by every segment, the last +//! describes a distinct value per segment: +//! +//! - ``cuda::args::constant{}`` for a value fixed at compile time. ``N`` is both the value and its bound. +//! - ``cuda::args::immediate{value}`` for a single value known on the host at the call. +//! - ``cuda::args::deferred{pointer}`` for a single value read in stream order, for example one produced on the +//! device by a preceding launch. +//! - ``cuda::args::deferred_sequence{iterator}`` for a distinct value per segment, also read in stream order. +//! +//! **How it is bounded.** A bound lets the algorithm reason about a value it does not know exactly: +//! +//! - A **compile-time** bound, ``cuda::args::bounds()``, may accompany ``immediate``, ``deferred``, or +//! ``deferred_sequence`` (a ``constant`` is already its own bound). The kernel specializes on this range and uses +//! it to size temporary storage (see *Choosing argument bounds*), so prefer the tightest range you can prove. +//! - A **runtime** bound, ``cuda::args::bounds(lo, hi)``, may accompany ``deferred`` and ``deferred_sequence`` when +//! the range is only known at runtime. A compile-time and a runtime bound can be combined, and the effective range +//! is then their intersection. +//! +//! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. The kernel specializes +//! on their upper bound, so each carries a compile-time bound (a ``constant`` is its own, the other forms take an +//! explicit ``cuda::args::bounds()``). ``num_segments`` is a single value, supplied as ``constant``, +//! ``immediate``, or ``deferred``, never as a per-segment sequence. +//! +//! .. code-block:: c++ +//! +//! // segment_sizes (k is analogous): +//! cuda::args::constant<256>{}; // fixed at compile time +//! cuda::args::immediate{n, cuda::args::bounds<1, 1024>()}; // host value, at most 1024 +//! cuda::args::deferred_sequence{d_sizes, cuda::args::bounds<1, 1024>()}; // per-segment, each at most 1024 +//! +//! // a single value produced on the device (e.g. num_segments), with a static cap and a tighter runtime cap: +//! cuda::args::deferred{d_count, cuda::args::bounds<0, 4096>(), cuda::args::bounds(0, n)}; //! //! Choosing argument bounds //! +++++++++++++++++++++++++++++++++++++++++++++ @@ -198,7 +226,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! The *statically-known maximum* segment size (the upper bound of the ``segment_sizes`` annotation) must be small //! enough that such a block fits within the shared-memory limit. Both uniform (fixed) and variable segment sizes are //! supported as long as this maximum is honored. -//! - **Uniform number of segments.** ``num_segments`` must be a single (host-resolved) value. +//! - **Uniform number of segments.** ``num_segments`` must be a single value, never a per-segment sequence. //! - **Explicit opt-out required for the output guarantees.** The deterministic, stable-sorted default contract //! described in *Determinism, tie-breaking, and output ordering* below (and in :ref:`cub-topk-requirements`) is not //! yet implemented. Like :cpp:struct:`cub::DeviceTopK`, the caller must currently request non-deterministic, @@ -585,16 +613,7 @@ struct DeviceBatchedTopK return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { return detail::dispatch_batched_topk( - storage, - bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k, - num_segments, - env); + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env); }); } @@ -690,16 +709,7 @@ struct DeviceBatchedTopK return detail::dispatch_with_env( env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { return detail::dispatch_batched_topk( - storage, - bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k, - num_segments, - env); + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env); }); } }; diff --git a/cub/examples/device/example_device_batched_topk_pairs.cu b/cub/examples/device/example_device_batched_topk_pairs.cu index aa7adbdd7a89..7dc15e1cc46e 100644 --- a/cub/examples/device/example_device_batched_topk_pairs.cu +++ b/cub/examples/device/example_device_batched_topk_pairs.cu @@ -97,9 +97,9 @@ int main(int argc, char** argv) auto env = cuda::std::execution::env{cuda::stream_ref{stream}, requirements}; // Annotate the arguments: runtime segment size and K, each with a compile-time upper bound. - auto segment_sizes = cuda::args::immediate{segment_size, cuda::args::bounds()}; - auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; - auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; + auto segment_sizes = cuda::args::immediate{segment_size, cuda::args::bounds()}; + auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; + auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; size_t temp_storage_bytes = 0; CubDebugExit(DeviceBatchedTopK::MaxPairs( diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu index 3365715236c5..4fc1122201f1 100644 --- a/cub/test/catch2_test_device_batched_topk_api.cu +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -157,16 +157,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ size_t temp_storage_bytes = 0; auto error = cub::DeviceBatchedTopK::MaxPairs( - nullptr, - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k_arg, - num_segs, - env); + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k_arg, num_segs, env); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); if (error == cudaSuccess) { @@ -240,16 +231,7 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ size_t temp_storage_bytes = 0; auto error = cub::DeviceBatchedTopK::MinPairs( - nullptr, - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k_arg, - num_segs, - env); + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k_arg, num_segs, env); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); if (error == cudaSuccess) { diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 542cdc4eba45..fc27f26920d4 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -52,26 +52,12 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( if constexpr (SelectDirection == cub::detail::topk::select::max) { return cub::DeviceBatchedTopK::MaxKeys( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - segment_sizes, - k, - num_segments, - env); + d_temp_storage, temp_storage_bytes, d_key_segments_it, d_key_segments_out_it, segment_sizes, k, num_segments, env); } else { return cub::DeviceBatchedTopK::MinKeys( - d_temp_storage, - temp_storage_bytes, - d_key_segments_it, - d_key_segments_out_it, - segment_sizes, - k, - num_segments, - env); + d_temp_storage, temp_storage_bytes, d_key_segments_it, d_key_segments_out_it, segment_sizes, k, num_segments, env); } } From fbd66d787347540bed4f6e89c2be1addac7d5b1b Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 16 Jun 2026 22:51:02 -0700 Subject: [PATCH 06/18] drops standalone examples for follow-up pr --- .../example_device_batched_topk_keys.cu | 135 --------------- .../example_device_batched_topk_pairs.cu | 155 ------------------ 2 files changed, 290 deletions(-) delete mode 100644 cub/examples/device/example_device_batched_topk_keys.cu delete mode 100644 cub/examples/device/example_device_batched_topk_pairs.cu diff --git a/cub/examples/device/example_device_batched_topk_keys.cu b/cub/examples/device/example_device_batched_topk_keys.cu deleted file mode 100644 index 6bbf10054bbb..000000000000 --- a/cub/examples/device/example_device_batched_topk_keys.cu +++ /dev/null @@ -1,135 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -//! Simple example of cub::DeviceBatchedTopK::MaxKeys(). -//! For each (small) segment, find the K largest float keys. This example also demonstrates the argument annotation -//! framework: the (small) segment size and K are passed as compile-time `cuda::args::constant<>` values, while -//! the number of segments is a runtime `cuda::args::immediate`. -//! -//! To compile using the command line: -//! nvcc -arch=sm_XX example_device_batched_topk_keys.cu -I../.. -lcudart -O3 - -// Ensure printing of CUDA runtime errors to console -#define CUB_STDERR - -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "../../test/test_util.h" - -using namespace cub; - -// Whether to display input/output to console -bool g_verbose = false; - -// Compile-time, statically-known (small) segment size and K. Passing them as compile-time constants lets -// cub::DeviceBatchedTopK specialize the kernel for a single thread block per segment. -static constexpr int segment_size = 256; -static constexpr int k = 8; - -int main(int argc, char** argv) -{ - int num_segments = 4000; // runtime number of segments - - // Initialize command line - CommandLineArgs args(argc, argv); - g_verbose = args.CheckCmdLineFlag("v"); - args.GetCmdLineArgument("num-segments", num_segments); - - if (args.CheckCmdLineFlag("help")) - { - std::cout << "Usage: " << argv[0] << " [--num-segments=] [--device=] [--v]\n"; - exit(0); - } - - CubDebugExit(args.DeviceInit()); - - const int num_items = num_segments * segment_size; - std::cout << "cub::DeviceBatchedTopK::MaxKeys() finds the " << k << " largest items in each of " << num_segments - << " segments of " << segment_size << " items (" << sizeof(float) << "-byte keys)\n"; - - // Initialize host input - thrust::host_vector h_keys_in(num_items); - for (int i = 0; i < num_items; ++i) - { - RandomBits(h_keys_in[i]); - } - - // Compute the reference solution: largest K of each segment, sorted descending - std::vector reference(static_cast(num_segments) * k); - for (int s = 0; s < num_segments; ++s) - { - std::vector seg(h_keys_in.begin() + s * segment_size, h_keys_in.begin() + (s + 1) * segment_size); - std::partial_sort(seg.begin(), seg.begin() + k, seg.end(), std::greater{}); - std::copy(seg.begin(), seg.begin() + k, reference.begin() + static_cast(s) * k); - } - - // Allocate device arrays - thrust::device_vector d_keys_in_buffer(h_keys_in); - thrust::device_vector d_keys_out_buffer(static_cast(num_segments) * k, thrust::no_init); - - // Build per-segment iterators: d_keys_in[s] points to the start of segment s. - auto d_keys_in = cuda::make_strided_iterator( - cuda::make_counting_iterator(thrust::raw_pointer_cast(d_keys_in_buffer.data())), segment_size); - auto d_keys_out = - cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(d_keys_out_buffer.data())), k); - - // Specify that we do not require a specific output order and do not require deterministic results - auto requirements = - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted); - cudaStream_t stream = nullptr; - CubDebugExit(cudaStreamCreate(&stream)); - auto env = cuda::std::execution::env{cuda::stream_ref{stream}, requirements}; - - // Annotate the arguments: segment size and K are compile-time constants; the number of segments is a runtime value. - auto segment_sizes = cuda::args::constant{}; - auto k_arg = cuda::args::constant{}; - auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; - - // Query temporary storage requirements - size_t temp_storage_bytes = 0; - CubDebugExit(DeviceBatchedTopK::MaxKeys( - nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs_arg, env)); - - // Allocate temporary storage - thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - void* d_temp_storage = thrust::raw_pointer_cast(temp_storage.data()); - - // Run the segmented top-k algorithm - CubDebugExit(DeviceBatchedTopK::MaxKeys( - d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs_arg, env)); - - // Check for correctness: the per-segment output is unordered, so sort each output segment descending before - // comparing against the (descending) reference. - thrust::host_vector h_keys_out(d_keys_out_buffer); - int compare = 0; - for (int s = 0; s < num_segments && compare == 0; ++s) - { - std::sort(h_keys_out.begin() + s * k, h_keys_out.begin() + (s + 1) * k, std::greater{}); - compare = CompareResults( - reference.data() + static_cast(s) * k, - thrust::raw_pointer_cast(h_keys_out.data()) + static_cast(s) * k, - k, - g_verbose); - } - AssertEquals(0, compare); - - CubDebugExit(cudaStreamDestroy(stream)); - std::cout << "\n\n"; - return 0; -} diff --git a/cub/examples/device/example_device_batched_topk_pairs.cu b/cub/examples/device/example_device_batched_topk_pairs.cu deleted file mode 100644 index 7dc15e1cc46e..000000000000 --- a/cub/examples/device/example_device_batched_topk_pairs.cu +++ /dev/null @@ -1,155 +0,0 @@ -// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception - -//! Simple example of cub::DeviceBatchedTopK::MaxPairs(). -//! For each (small) segment, find the K largest float keys and gather their associated values. This example -//! demonstrates the argument annotation framework with *runtime* values that carry a *compile-time* upper bound: -//! the segment size and K are passed as `cuda::args::immediate{value, cuda::args::bounds<1, MAX>()}`. -//! -//! To compile using the command line: -//! nvcc -arch=sm_XX example_device_batched_topk_pairs.cu -I../.. -lcudart -O3 - -// Ensure printing of CUDA runtime errors to console -#define CUB_STDERR - -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include "../../test/test_util.h" - -using namespace cub; - -bool g_verbose = false; - -// Compile-time upper bounds on the (small) segment size and K. The actual values are provided at runtime but are -// guaranteed not to exceed these bounds, which lets cub::DeviceBatchedTopK specialize for a single block per segment. -using offset_t = cuda::std::int64_t; -static constexpr offset_t max_seg_size = 1024; -static constexpr offset_t max_k = 32; - -int main(int argc, char** argv) -{ - offset_t segment_size = 256; - offset_t k = 8; - int num_segments = 2000; - - CommandLineArgs args(argc, argv); - g_verbose = args.CheckCmdLineFlag("v"); - args.GetCmdLineArgument("segment-size", segment_size); - args.GetCmdLineArgument("k", k); - args.GetCmdLineArgument("num-segments", num_segments); - - if (args.CheckCmdLineFlag("help")) - { - std::cout << "Usage: " << argv[0] - << " [--segment-size=] [--k=] [--num-segments=] [--device=] [--v]\n"; - exit(0); - } - CubDebugExit(args.DeviceInit()); - - segment_size = std::min(segment_size, max_seg_size); - k = std::min(k, segment_size); - const offset_t num_items = static_cast(num_segments) * segment_size; - - std::cout << "cub::DeviceBatchedTopK::MaxPairs() finds the " << k << " largest items in each of " << num_segments - << " segments of " << segment_size << " items\n"; - - // Initialize host input keys; values are the global item indices. - thrust::host_vector h_keys_in(num_items); - for (offset_t i = 0; i < num_items; ++i) - { - RandomBits(h_keys_in[i]); - } - - thrust::device_vector d_keys_in_buffer(h_keys_in); - thrust::device_vector d_keys_out_buffer(static_cast(num_segments) * k, thrust::no_init); - thrust::device_vector d_values_out_buffer(static_cast(num_segments) * k, thrust::no_init); - - // Per-segment iterators. Input values are the global indices [0, num_items) via a counting iterator. - auto d_keys_in = cuda::make_strided_iterator( - cuda::make_counting_iterator(thrust::raw_pointer_cast(d_keys_in_buffer.data())), segment_size); - auto d_keys_out = - cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(d_keys_out_buffer.data())), k); - auto d_values_in = - cuda::make_strided_iterator(cuda::make_counting_iterator(cuda::make_counting_iterator(offset_t{0})), segment_size); - auto d_values_out = - cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(d_values_out_buffer.data())), k); - - auto requirements = - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted); - cudaStream_t stream = nullptr; - CubDebugExit(cudaStreamCreate(&stream)); - auto env = cuda::std::execution::env{cuda::stream_ref{stream}, requirements}; - - // Annotate the arguments: runtime segment size and K, each with a compile-time upper bound. - auto segment_sizes = cuda::args::immediate{segment_size, cuda::args::bounds()}; - auto k_arg = cuda::args::immediate{k, cuda::args::bounds()}; - auto num_segs_arg = cuda::args::immediate{static_cast(num_segments)}; - - size_t temp_storage_bytes = 0; - CubDebugExit(DeviceBatchedTopK::MaxPairs( - nullptr, - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k_arg, - num_segs_arg, - env)); - - thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - CubDebugExit(DeviceBatchedTopK::MaxPairs( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k_arg, - num_segs_arg, - env)); - - // Validate: (1) each output value indexes back to the matching output key, and (2) the selected keys are the K - // largest of their segment. - thrust::host_vector h_keys_out(d_keys_out_buffer); - thrust::host_vector h_values_out(d_values_out_buffer); - bool ok = true; - for (offset_t s = 0; s < num_segments && ok; ++s) - { - // (1) value/key association - for (offset_t j = 0; j < k; ++j) - { - const offset_t global_idx = h_values_out[s * k + j]; - ok = ok && global_idx >= 0 && global_idx < num_items && h_keys_in[global_idx] == h_keys_out[s * k + j]; - } - // (2) key correctness against a host reference (top-k largest, order-independent) - std::vector seg(h_keys_in.begin() + s * segment_size, h_keys_in.begin() + (s + 1) * segment_size); - std::partial_sort(seg.begin(), seg.begin() + k, seg.end(), std::greater{}); - std::vector res(h_keys_out.begin() + s * k, h_keys_out.begin() + (s + 1) * k); - std::sort(res.begin(), res.end(), std::greater{}); - ok = ok && std::equal(seg.begin(), seg.begin() + k, res.begin()); - } - AssertEquals(true, ok); - - CubDebugExit(cudaStreamDestroy(stream)); - std::cout << "\n\n"; - return 0; -} From b6c811312426573ab8c4f6a4324fba225487de47 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Wed, 17 Jun 2026 08:10:26 -0700 Subject: [PATCH 07/18] addresses review comments --- cub/cub/device/device_batched_topk.cuh | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index ee0038f7cdc5..fd85f9c039f8 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -111,8 +111,8 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( // --------------------------------------------------------------------------- // Resolve the (optionally tuned) policy selector from the environment. // --------------------------------------------------------------------------- - using key_t = it_value_t>; - using value_t = it_value_t>; + using key_t = cub::detail::it_value_t>; + using value_t = cub::detail::it_value_t>; using default_policy_selector_t = batched_topk:: policy_selector_from_types::highest>; using tuning_env_t = @@ -166,7 +166,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! Argument annotation framework //! +++++++++++++++++++++++++++++++++++++++++++++ //! -//! The parameters ``segment_sizes``, ``k``, and ``num_segments`` are passed as **annotated arguments** from +//! The parameters ``segment_sizes``, ``k``, and ``num_segments`` can be passed as **annotated arguments** from //! ``cuda::args``. An annotation tells the algorithm everything you know about a parameter: where its value comes //! from and how tightly it is bounded. The more you can tell the algorithm, and the more precisely (a //! compile-time constant rather than a runtime value, a tight bound rather than a loose one), the more it can @@ -177,8 +177,8 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! //! - ``cuda::args::constant{}`` for a value fixed at compile time. ``N`` is both the value and its bound. //! - ``cuda::args::immediate{value}`` for a single value known on the host at the call. -//! - ``cuda::args::deferred{pointer}`` for a single value read in stream order, for example one produced on the -//! device by a preceding launch. +//! - ``cuda::args::deferred{iterator}`` for a single value read in stream order through a pointer or iterator, for +//! example one produced on the device by a preceding launch. //! - ``cuda::args::deferred_sequence{iterator}`` for a distinct value per segment, also read in stream order. //! //! **How it is bounded.** A bound lets the algorithm reason about a value it does not know exactly: @@ -187,13 +187,13 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! ``deferred_sequence`` (a ``constant`` is already its own bound). The kernel specializes on this range and uses //! it to size temporary storage (see *Choosing argument bounds*), so prefer the tightest range you can prove. //! - A **runtime** bound, ``cuda::args::bounds(lo, hi)``, may accompany ``deferred`` and ``deferred_sequence`` when -//! the range is only known at runtime. A compile-time and a runtime bound can be combined, and the effective range -//! is then their intersection. +//! the range is only known at runtime. When combined with a compile-time bound, the runtime bound must be at least +//! as narrow, lying within the compile-time range and only tightening it further. //! //! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. The kernel specializes //! on their upper bound, so each carries a compile-time bound (a ``constant`` is its own, the other forms take an -//! explicit ``cuda::args::bounds()``). ``num_segments`` is a single value, supplied as ``constant``, -//! ``immediate``, or ``deferred``, never as a per-segment sequence. +//! explicit ``cuda::args::bounds()``). ``num_segments`` must be a single value known on the host, supplied +//! as ``constant`` or ``immediate``. //! //! .. code-block:: c++ //! From 8111d1b97091640963d5ae573e60477fc02bd275 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Mon, 22 Jun 2026 00:46:33 -0700 Subject: [PATCH 08/18] refines docs --- cub/cub/device/device_batched_topk.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index fd85f9c039f8..f7c5b5712296 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -202,8 +202,8 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! cuda::args::immediate{n, cuda::args::bounds<1, 1024>()}; // host value, at most 1024 //! cuda::args::deferred_sequence{d_sizes, cuda::args::bounds<1, 1024>()}; // per-segment, each at most 1024 //! -//! // a single value produced on the device (e.g. num_segments), with a static cap and a tighter runtime cap: -//! cuda::args::deferred{d_count, cuda::args::bounds<0, 4096>(), cuda::args::bounds(0, n)}; +//! // a uniform segment size produced on the device, capped at compile time and narrowed at runtime: +//! cuda::args::deferred{d_size, cuda::args::bounds<1, 1024>(), cuda::args::bounds(1, runtime_max)}; //! //! Choosing argument bounds //! +++++++++++++++++++++++++++++++++++++++++++++ From 2f7f695d2a1c66f54bb5c345153d8093fcb4dc95 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Mon, 22 Jun 2026 21:12:05 -0700 Subject: [PATCH 09/18] introduces the both-or-nothing rule for determinism and tie-break --- cub/cub/device/device_batched_topk.cuh | 98 ++++++++----- cub/cub/device/device_topk.cuh | 11 +- .../catch2_test_device_batched_topk_api.cu | 24 ++-- ...catch2_test_device_batched_topk_env_api.cu | 16 ++- .../catch2_test_device_segmented_topk_keys.cu | 4 +- ...catch2_test_device_segmented_topk_pairs.cu | 4 +- ...t_device_batched_topk_requirements_fail.cu | 68 +++++++++ .../cub/api_docs/device_topk_requirements.rst | 134 +++++++++++------- 8 files changed, 260 insertions(+), 99 deletions(-) create mode 100644 cub/test/test_device_batched_topk_requirements_fail.cu diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index f7c5b5712296..0b379d410501 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -72,41 +72,68 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( EnvT env) { // --------------------------------------------------------------------------- - // Execution requirements (mirrors cub::DeviceTopK): non-deterministic + unsorted output only. + // Execution requirements. + // + // Two orthogonal concerns govern the result: *which* items are selected (possibly refined by + // a tie-break preference) and the order in which they are written (output ordering). The committed default contract + // is the most reproducible behavior (determinism::gpu_to_gpu + tie_break::prefer_smaller_index + + // output_ordering::stable_sorted). Three rules are validated here: + // 1. determinism and tie_break must be acknowledged together (both specified, or both omitted default) + // 2. an explicit tie_break of prefer_smaller_index / prefer_larger_index fully pins the result set across GPUs and + // therefore requires determinism::gpu_to_gpu (it cannot be paired with run_to_run or not_guaranteed) + // 3. this initial API surface only implements the fully opted-out configuration (non-deterministic, unsorted). // --------------------------------------------------------------------------- static_assert(!::cuda::std::execution::__queryable_with, - "Determinism should be used inside requires to have an effect."); + "Determinism should be used inside cuda::execution::require to have an effect."); + static_assert(!::cuda::std::execution::__queryable_with, + "Tie-break should be used inside cuda::execution::require to have an effect."); using requirements_t = ::cuda::std::execution:: __query_result_or_t>; + + constexpr bool determinism_specified = + ::cuda::std::execution::__queryable_with; + constexpr bool tie_break_specified = + ::cuda::std::execution::__queryable_with; + using requested_determinism_t = ::cuda::std::execution::__query_result_or_t; + ::cuda::execution::determinism::gpu_to_gpu_t>; + using requested_tie_break_t = + ::cuda::std::execution::__query_result_or_t; using requested_order_t = ::cuda::std::execution::__query_result_or_t; - static_assert(::cuda::std::is_same_v - && ::cuda::std::is_same_v, - "cub::DeviceBatchedTopK only supports non-deterministic, unsorted output. Acknowledge this by " - "passing cuda::execution::require(cuda::execution::determinism::not_guaranteed, " - "cuda::execution::output_ordering::unsorted) in the environment."); + ::cuda::execution::output_ordering::stable_sorted_t>; - // A tie-break requirement only constrains *which* of the elements that compare equal at the K-th position are - // selected (per segment); it is meaningless without determinism, so it must be paired with a deterministic - // requirement (run_to_run / gpu_to_gpu). Deterministic execution and tie-breaking are not implemented yet; this only - // validates the requirement combination so the eventual behavior is wired up. - using requested_tie_break_t = - ::cuda::std::execution::__query_result_or_t; - static_assert( + constexpr bool determinism_and_tie_break_paired = (determinism_specified == tie_break_specified); + + // Encodes rule 2 as the implication "a concrete tie-break requires gpu_to_gpu". The expression is the "or" form + // and satisfied in the two cases that are allowed: the tie-break is unspecified or the determinism is already + // gpu_to_gpu (which accepts any tie-break). + constexpr bool tie_break_compatible_with_determinism = ::cuda::std::is_same_v - || !::cuda::std::is_same_v, - "cub::DeviceBatchedTopK: a tie_break requirement (cuda::execution::tie_break::prefer_smaller_index or " - "prefer_larger_index) requires a deterministic execution requirement " - "(cuda::execution::determinism::run_to_run or gpu_to_gpu); it cannot be combined with " - "cuda::execution::determinism::not_guaranteed."); + || ::cuda::std::is_same_v; + constexpr bool is_non_deterministic_unsorted = + ::cuda::std::is_same_v + && ::cuda::std::is_same_v; + + static_assert(determinism_and_tie_break_paired, + "cub::DeviceBatchedTopK: determinism and tie_break requirements must be acknowledged together. Either " + "omit both to accept the defaults (cuda::execution::determinism::gpu_to_gpu and " + "cuda::execution::tie_break::prefer_smaller_index), or pass both explicitly inside " + "cuda::execution::require(...)."); + static_assert(!determinism_and_tie_break_paired || tie_break_compatible_with_determinism, + "cub::DeviceBatchedTopK: a tie_break of cuda::execution::tie_break::prefer_smaller_index or " + "prefer_larger_index pins the result set across GPUs and therefore requires " + "cuda::execution::determinism::gpu_to_gpu (it cannot be combined with run_to_run or not_guaranteed)."); + static_assert( + !determinism_and_tie_break_paired || !tie_break_compatible_with_determinism || is_non_deterministic_unsorted, + "cub::DeviceBatchedTopK currently only implements non-deterministic, unsorted output. Request it " + "explicitly with cuda::execution::require(cuda::execution::determinism::not_guaranteed, " + "cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)."); // --------------------------------------------------------------------------- // Resolve the (optionally tuned) policy selector from the environment. @@ -229,9 +256,10 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! - **Uniform number of segments.** ``num_segments`` must be a single value, never a per-segment sequence. //! - **Explicit opt-out required for the output guarantees.** The deterministic, stable-sorted default contract //! described in *Determinism, tie-breaking, and output ordering* below (and in :ref:`cub-topk-requirements`) is not -//! yet implemented. Like :cpp:struct:`cub::DeviceTopK`, the caller must currently request non-deterministic, -//! unsorted output explicitly by passing ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, -//! cuda::execution::output_ordering::unsorted)`` in the environment. +//! yet implemented. The caller must currently request non-deterministic, unsorted output explicitly by passing +//! ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, +//! cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)`` in the environment +//! (``determinism`` and ``tie_break`` must always be specified together). //! //! Determinism, tie-breaking, and output ordering //! +++++++++++++++++++++++++++++++++++++++++++++++ @@ -240,10 +268,12 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! requirements: *which* items are selected per segment (``cuda::execution::determinism``, optionally refined by //! ``cuda::execution::tie_break``) and the order in which they are written (``cuda::execution::output_ordering``). //! When the caller does not opt out, the committed default is the most reproducible behavior: deterministic results -//! (``cuda::execution::determinism::run_to_run``), ties resolved toward the smaller (lower) source index +//! (``cuda::execution::determinism::gpu_to_gpu``), ties resolved toward the smaller (lower) source index //! (``cuda::execution::tie_break::prefer_smaller_index``), and stable-sorted output //! (``cuda::execution::output_ordering::stable_sorted``). Callers opt *out* of these guarantees to obtain faster -//! implementations. +//! implementations. ``determinism`` and ``tie_break`` must always be specified together, or both omitted to take the +//! default. A specified ``tie_break`` of ``prefer_smaller_index`` or ``prefer_larger_index`` requires +//! ``determinism::gpu_to_gpu``. //! //! See :ref:`cub-topk-requirements` for the full requirement model, worked examples, and guidance on choosing //! requirements. @@ -252,10 +282,10 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! //! **Current support.** This release only implements the fully opted-out configuration, which must be requested //! explicitly: ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, -//! cuda::execution::output_ordering::unsorted)``. Any other combination (including an empty, no-requirement -//! environment) is rejected at compile time. In this configuration the per-segment output is unordered and may be -//! non-deterministic: if multiple items tie at the K-th position, the subset of tied elements returned is not -//! uniquely defined and may vary between runs. +//! cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)``. Any other combination +//! (including an empty, no-requirement environment) is rejected at compile time. In this configuration the +//! per-segment output is unordered and may be non-deterministic: if multiple items tie at the K-th position, the +//! subset of tied elements returned is not uniquely defined and may vary between runs. //! //! Usage Considerations //! ++++++++++++++++++++++++++ @@ -321,8 +351,8 @@ struct DeviceBatchedTopK //! //! @param[in] env //! @rst - //! **[optional]** Execution environment. Must require `determinism::not_guaranteed` and - //! `output_ordering::unsorted`. + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, + //! `tie_break::unspecified`, and `output_ordering::unsorted`. //! @endrst template ; + ::cuda::execution::determinism::gpu_to_gpu_t>; using requested_order_t = ::cuda::std::execution::__query_result_or_t; + ::cuda::execution::output_ordering::stable_sorted_t>; constexpr auto is_determinism_not_guaranteed = ::cuda::std::is_same_v; constexpr auto is_output_order_unsorted = @@ -84,6 +84,11 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_topk( "cub::DeviceTopK only supports the case where determinism is not guaranteed and output order is " "unsorted."); + // TODO (elstehle): align requirement validation with cub::DeviceBatchedTopK in CCCL 4.0. cub::DeviceTopK does not + // yet inspect cuda::execution::tie_break, so it still accepts requirement combinations that cub::DeviceBatchedTopK + // rejects. It should enforce that determinism and tie_break are requested together (or both omitted to take the + // default) and that an explicit tie_break requires cuda::execution::determinism::gpu_to_gpu. + // Query relevant properties from the environment auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); @@ -167,7 +172,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_topk_hub( //! The result of ``DeviceTopK`` is governed by two orthogonal execution requirements: *which* items are selected //! (``cuda::execution::determinism``, optionally refined by ``cuda::execution::tie_break``) and the order in which //! they are written (``cuda::execution::output_ordering``). When the caller does not opt out, the committed default -//! is the most reproducible behavior: deterministic results (``cuda::execution::determinism::run_to_run``), ties +//! is the most reproducible behavior: deterministic results (``cuda::execution::determinism::gpu_to_gpu``), ties //! resolved toward the smaller (lower) source index (``cuda::execution::tie_break::prefer_smaller_index``), and //! stable-sorted output (``cuda::execution::output_ordering::stable_sorted``). Callers opt *out* of these guarantees //! to obtain faster implementations. diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu index 4fc1122201f1..9061f55001b4 100644 --- a/cub/test/catch2_test_device_batched_topk_api.cu +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -45,8 +45,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_t auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; // Top-k output is unordered and may be non-deterministic; this must be acknowledged via the environment. - auto env = cuda::std::execution::env{ - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; // Query temporary storage requirements size_t temp_storage_bytes = 0; @@ -98,8 +100,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys temp-storage API example", "[batched_t auto segment_sizes = cuda::args::constant{}; auto k_arg = cuda::args::constant{}; auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto env = cuda::std::execution::env{ - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; size_t temp_storage_bytes = 0; auto error = cub::DeviceBatchedTopK::MinKeys( @@ -152,8 +156,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ auto segment_sizes = cuda::args::constant{}; auto k_arg = cuda::args::constant{}; auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto env = cuda::std::execution::env{ - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; size_t temp_storage_bytes = 0; auto error = cub::DeviceBatchedTopK::MaxPairs( @@ -226,8 +232,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ auto segment_sizes = cuda::args::constant{}; auto k_arg = cuda::args::constant{}; auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto env = cuda::std::execution::env{ - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; size_t temp_storage_bytes = 0; auto error = cub::DeviceBatchedTopK::MinPairs( diff --git a/cub/test/catch2_test_device_batched_topk_env_api.cu b/cub/test/catch2_test_device_batched_topk_env_api.cu index caa21e802cc8..0fa51808e166 100644 --- a/cub/test/catch2_test_device_batched_topk_env_api.cu +++ b/cub/test/catch2_test_device_batched_topk_env_api.cu @@ -41,7 +41,9 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys env-alloc example", "[batched_topk][de cuda::stream stream{cuda::devices[0]}; auto env = cuda::std::execution::env{ - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted), + cuda::execution::require(cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted), cuda::stream_ref{stream}}; // The env-based overload allocates and frees the temporary storage internally. @@ -82,7 +84,9 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys env-alloc example", "[batched_topk][de cuda::stream stream{cuda::devices[0]}; auto env = cuda::std::execution::env{ - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted), + cuda::execution::require(cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted), cuda::stream_ref{stream}}; auto error = cub::DeviceBatchedTopK::MinKeys( @@ -126,7 +130,9 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][d cuda::stream stream{cuda::devices[0]}; auto env = cuda::std::execution::env{ - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted), + cuda::execution::require(cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted), cuda::stream_ref{stream}}; auto error = cub::DeviceBatchedTopK::MaxPairs( @@ -191,7 +197,9 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][d cuda::stream stream{cuda::devices[0]}; auto env = cuda::std::execution::env{ - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted), + cuda::execution::require(cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted), cuda::stream_ref{stream}}; auto error = cub::DeviceBatchedTopK::MinPairs( diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index fc27f26920d4..86c1a00a0267 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -48,7 +48,9 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_keys( { auto env = cuda::std::execution::env{ cuda::stream_ref{stream}, - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + cuda::execution::require(cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; if constexpr (SelectDirection == cub::detail::topk::select::max) { return cub::DeviceBatchedTopK::MaxKeys( diff --git a/cub/test/catch2_test_device_segmented_topk_pairs.cu b/cub/test/catch2_test_device_segmented_topk_pairs.cu index 2143e479e3e3..4b684f5d5d07 100644 --- a/cub/test/catch2_test_device_segmented_topk_pairs.cu +++ b/cub/test/catch2_test_device_segmented_topk_pairs.cu @@ -72,7 +72,9 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk_pairs( { auto env = cuda::std::execution::env{ cuda::stream_ref{stream}, - cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)}; + cuda::execution::require(cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; if constexpr (SelectDirection == cub::detail::topk::select::max) { return cub::DeviceBatchedTopK::MaxPairs( diff --git a/cub/test/test_device_batched_topk_requirements_fail.cu b/cub/test/test_device_batched_topk_requirements_fail.cu new file mode 100644 index 000000000000..c9a00e9d31d3 --- /dev/null +++ b/cub/test/test_device_batched_topk_requirements_fail.cu @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +// %PARAM% TEST_ERR err 0:1:2:3:4:5 + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +// Verifies that cub::DeviceBatchedTopK rejects, at compile time, the determinism and tie_break requirement +// combinations that the public contract marks as ill-formed (see docs/cub/api_docs/device_topk_requirements.rst): +// * determinism and tie_break must be acknowledged together, both specified or both omitted to take the default +// * an explicit tie_break of prefer_smaller_index or prefer_larger_index pins the result set across GPUs and so +// requires determinism::gpu_to_gpu (it cannot be paired with not_guaranteed or run_to_run) +// Each variant exercises one rejected cell of that contract table. + +int main() +{ + namespace ex = cuda::execution; + + // Per-segment key iterators (iterator-of-iterators). The keys-only entry point ignores value iterators. + int** d_keys_in = nullptr; + int** d_keys_out = nullptr; + auto segment_sizes = cuda::args::constant<8>{}; + auto k_arg = cuda::args::constant<3>{}; + auto num_segments = cuda::args::immediate{cuda::std::int64_t{2}}; + +#if TEST_ERR == 0 // determinism specified without a paired tie_break + auto requirements = ex::require(ex::determinism::not_guaranteed, ex::output_ordering::unsorted); + // expected-error-0 {{"must be acknowledged together"}} +#elif TEST_ERR == 1 // tie_break specified without a paired determinism + auto requirements = ex::require(ex::tie_break::prefer_smaller_index, ex::output_ordering::unsorted); + // expected-error-1 {{"must be acknowledged together"}} +#elif TEST_ERR == 2 // explicit tie_break with not_guaranteed + auto requirements = + ex::require(ex::determinism::not_guaranteed, ex::tie_break::prefer_smaller_index, ex::output_ordering::unsorted); + // expected-error-2 {{"pins the result set across GPUs and therefore requires"}} +#elif TEST_ERR == 3 // explicit tie_break with not_guaranteed + auto requirements = + ex::require(ex::determinism::not_guaranteed, ex::tie_break::prefer_larger_index, ex::output_ordering::unsorted); + // expected-error-3 {{"pins the result set across GPUs and therefore requires"}} +#elif TEST_ERR == 4 // explicit tie_break with run_to_run + auto requirements = + ex::require(ex::determinism::run_to_run, ex::tie_break::prefer_smaller_index, ex::output_ordering::unsorted); + // expected-error-4 {{"pins the result set across GPUs and therefore requires"}} +#elif TEST_ERR == 5 // explicit tie_break with run_to_run + auto requirements = + ex::require(ex::determinism::run_to_run, ex::tie_break::prefer_larger_index, ex::output_ordering::unsorted); + // expected-error-5 {{"pins the result set across GPUs and therefore requires"}} +#endif + + auto env = cuda::std::execution::env{requirements}; + size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segments, env); + if (error != cudaSuccess) + { + std::cerr << "cub::DeviceBatchedTopK::MaxKeys failed with status: " << error << '\n'; + } +} diff --git a/docs/cub/api_docs/device_topk_requirements.rst b/docs/cub/api_docs/device_topk_requirements.rst index 9b02280fea4b..a3117f9cea25 100644 --- a/docs/cub/api_docs/device_topk_requirements.rst +++ b/docs/cub/api_docs/device_topk_requirements.rst @@ -45,27 +45,34 @@ Default behavior When you do **not** specify any of these requirements, the top-k algorithms provide their strongest reproducibility guarantees. The committed default contract is: -* ``cuda::execution::determinism::run_to_run`` for a deterministic result set, +* ``cuda::execution::determinism::gpu_to_gpu`` for a deterministic result set, * ``cuda::execution::tie_break::prefer_smaller_index`` to resolve ties at the selection boundary toward the smaller (lower) source index, * ``cuda::execution::output_ordering::stable_sorted`` to write output sorted by key, with equal keys ordered by source index. -In other words, by default you get the same items, in the same positions, run after run. You opt -**out** of these guarantees (by requiring weaker properties such as -``cuda::execution::determinism::not_guaranteed`` and ``cuda::execution::output_ordering::unsorted``) -to obtain faster implementations. +In other words, by default you get the same items, in the same positions, run after run and across +GPUs of the same architecture. You opt **out** of these guarantees (by requiring weaker properties +such as ``cuda::execution::determinism::not_guaranteed`` and +``cuda::execution::output_ordering::unsorted``) to obtain faster implementations. + +``determinism`` and ``tie_break`` are coupled. You specify **both** of them (inside a single +``cuda::execution::require(...)``) or **neither** (to take the default). A specified ``tie_break`` of +``prefer_smaller_index`` or ``prefer_larger_index`` pins the result set across GPUs and therefore +requires ``determinism::gpu_to_gpu``. See :ref:`cub-topk-set-membership` for the full table. .. note:: - **Current support.** This initial API surface only implements the fully opted-out configuration: + **Current support.** This initial API surface only implements the fully opted-out configuration. + For :cpp:struct:`cub::DeviceBatchedTopK` it must be requested **explicitly** as ``cuda::execution::require(cuda::execution::determinism::not_guaranteed, - cuda::execution::output_ordering::unsorted)``. That configuration must be requested - **explicitly**. The algorithms ``static_assert`` for any other combination (including an empty, - no-requirement environment), so the deterministic default described above cannot yet be exercised - in code. The deterministic, tie-broken, and (stable-)sorted modes documented here define the - committed long-term contract and will become available (including as the no-requirement default) - as those code paths land. + cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)`` + (:cpp:struct:`cub::DeviceTopK` has no tie-break dimension yet and omits the ``tie_break`` token). + The algorithms ``static_assert`` for any other combination (including an empty, no-requirement + environment), so the deterministic default described above cannot yet be exercised in code. The + deterministic, tie-broken, and (stable-)sorted modes documented here define the committed long-term + contract and will become available (including as the no-requirement default) as those code paths + land. Requirements reference ---------------------- @@ -84,18 +91,21 @@ Determinism (``cuda::execution::determinism``) may be returned. Enables the fastest implementations. * - ``run_to_run`` - The result set is identical across repeated invocations on the same GPU with the same input. - The tie-breaking policy is implementation-defined unless ``tie_break`` is also specified. + The tie-breaking policy is implementation-defined. Pinning a specific tie-break is not + available at this level and requires ``gpu_to_gpu``. * - ``gpu_to_gpu`` - - The result set is identical across different GPUs of the same architecture. When ``tie_break`` - is explicitly set, the result set is fully pinned: ``run_to_run`` and ``gpu_to_gpu`` yield the - same set for a given input. + - The result set is identical across different GPUs of the same architecture. This is the only + level that may be combined with an explicit ``tie_break`` (``prefer_smaller_index`` or + ``prefer_larger_index``), which then fully pins the result set for a given input. Tie-break (``cuda::execution::tie_break``) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ -Only meaningful together with ``determinism::run_to_run`` or ``determinism::gpu_to_gpu``. It has no -effect with ``not_guaranteed`` (and pairing ``tie_break`` with ``not_guaranteed`` is rejected at -compile time). +A specified ``tie_break`` of ``prefer_smaller_index`` or ``prefer_larger_index`` pins the result set +across GPUs, so it requires ``determinism::gpu_to_gpu``. Pairing it with ``run_to_run`` or +``not_guaranteed`` is rejected at compile time. ``determinism`` and ``tie_break`` must always be +specified together (or both omitted to take the default). Use ``tie_break::unspecified`` to leave the +boundary policy to the implementation, for example alongside ``not_guaranteed`` or ``run_to_run``. .. list-table:: :header-rows: 1 @@ -103,14 +113,15 @@ compile time). * - Value - Meaning - * - ``unspecified`` *(API default)* - - Any deterministic tie-break is acceptable, and the implementation chooses. - * - ``prefer_smaller_index`` + * - ``unspecified`` + - Any deterministic tie-break is acceptable, and the implementation chooses. Valid with any + determinism level (including ``not_guaranteed`` and ``run_to_run``). + * - ``prefer_smaller_index`` *(default)* - Among elements that compare equal at the boundary, prefer those with the **smaller** source - index. + index. Requires ``determinism::gpu_to_gpu``. * - ``prefer_larger_index`` - Among elements that compare equal at the boundary, prefer those with the **larger** source - index. + index. Requires ``determinism::gpu_to_gpu``. Output ordering (``cuda::execution::output_ordering``) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -143,34 +154,60 @@ execution environment alongside other properties such as a stream: auto env = cuda::std::execution::env{ cuda::execution::require( - cuda::execution::determinism::run_to_run, + cuda::execution::determinism::gpu_to_gpu, cuda::execution::tie_break::prefer_smaller_index, cuda::execution::output_ordering::sorted), stream_ref}; +.. _cub-topk-set-membership: + Which items are selected? ------------------------- -Determinism and tie-break together control **set membership**: +Determinism and tie-break together control **set membership**. They are always specified as a pair +(or both omitted to take the default). Rows below are the ``determinism`` requirement and columns are +the paired ``tie_break`` requirement. Cells marked *(compile error)* are rejected by a ``static_assert``. .. list-table:: :header-rows: 1 - :widths: 55 45 + :stub-columns: 1 + :widths: 22 26 26 26 - * - ``require(...)`` - - Which items are selected - * - ``determinism::not_guaranteed`` - - Non-deterministic among tied elements (fast path) - * - ``determinism::run_to_run`` + * - ``determinism`` + - ``tie_break::unspecified`` + - ``tie_break::prefer_smaller_index`` + - ``tie_break::prefer_larger_index`` + * - ``not_guaranteed`` + - Non-deterministic (fast path) + - *(compile error)* + - *(compile error)* + * - ``run_to_run`` + - Deterministic, implementation-defined tie-break + - *(compile error)* + - *(compile error)* + * - ``gpu_to_gpu`` - Deterministic, implementation-defined tie-break - * - ``determinism::run_to_run, tie_break::prefer_smaller_index`` - - Deterministic, ties resolved toward the smaller source index - * - ``determinism::run_to_run, tie_break::prefer_larger_index`` - - Deterministic, ties resolved toward the larger source index + - Deterministic, ties toward the **smaller** source index + - Deterministic, ties toward the **larger** source index -The same table applies when using ``determinism::gpu_to_gpu`` instead of ``run_to_run``. When an -explicit ``tie_break`` is set, both determinism levels pin the result set to the same items for a -given input. +Reading the table: + +* A specified ``tie_break`` of ``prefer_smaller_index`` or ``prefer_larger_index`` pins the result set + across GPUs, which is a ``gpu_to_gpu`` guarantee. Requesting it alongside ``not_guaranteed`` or + ``run_to_run`` is a compile error, because you must acknowledge the ``gpu_to_gpu`` determinism you + receive. +* With ``tie_break::unspecified`` the implementation chooses the boundary policy. ``run_to_run`` and + ``gpu_to_gpu`` then differ only in *scope*: identical results on the same GPU versus across GPUs of + the same architecture. +* Omitting **both** requirements selects the default (``gpu_to_gpu`` with ``prefer_smaller_index``), + which is the bottom-middle cell. + +.. note:: + + This determinism and tie_break pairing rule is currently enforced only by + :cpp:struct:`cub::DeviceBatchedTopK`. :cpp:struct:`cub::DeviceTopK` does not yet inspect + ``tie_break``, so it still accepts requirement combinations that ``cub::DeviceBatchedTopK`` rejects. + The same enforcement will be added to ``cub::DeviceTopK`` in the next major release of CCCL (4.0). Worked example: set membership x output ordering ------------------------------------------------- @@ -199,7 +236,8 @@ membership varies. - ``output_ordering::unsorted`` - ``output_ordering::sorted`` - ``output_ordering::stable_sorted`` - * - ``determinism::not_guaranteed`` + * - ``determinism::not_guaranteed,`` + ``tie_break::unspecified`` - | Run 1: ``[8@2, 10@0, 8@1]`` | Run 2: ``[8@3, 10@0, 8@1]`` | Different sets *and* orders @@ -209,8 +247,8 @@ membership varies. - | Run 1: ``[10@0, 8@1, 8@2]`` | Run 2: ``[10@0, 8@1, 8@3]`` | Different sets, equal keys in input order - * - ``determinism::run_to_run`` - (impl-defined tie-break) + * - ``determinism::run_to_run,`` + ``tie_break::unspecified`` - | Run 1: ``[8@3, 10@0, 8@1]`` | Run 2: ``[10@0, 8@1, 8@3]`` | Same set ``{10@0, 8@1, 8@3}``, order may vary @@ -220,7 +258,7 @@ membership varies. - | Run 1: ``[10@0, 8@1, 8@3]`` | Run 2: ``[10@0, 8@1, 8@3]`` | Same set, equal keys always in input order - * - ``determinism::run_to_run,`` + * - ``determinism::gpu_to_gpu,`` ``tie_break::prefer_smaller_index`` - | Run 1: ``[8@2, 10@0, 8@1]`` | Run 2: ``[10@0, 8@1, 8@2]`` @@ -231,7 +269,7 @@ membership varies. - | Run 1: ``[10@0, 8@1, 8@2]`` | Run 2: ``[10@0, 8@1, 8@2]`` | Same set, equal keys always in input order - * - ``determinism::run_to_run,`` + * - ``determinism::gpu_to_gpu,`` ``tie_break::prefer_larger_index`` - | Run 1: ``[8@3, 10@0, 8@2]`` | Run 2: ``[10@0, 8@2, 8@3]`` @@ -258,7 +296,7 @@ Reading the matrix: * - Set membership fixed, sorted but unstable among equal keys - ``run_to_run`` + ``sorted``: both runs start with ``10@0``, but ``8@1`` and ``8@3`` may swap * - Fully pinned: same set and same order - - ``run_to_run`` + ``tie_break::prefer_smaller_index`` + ``stable_sorted``: both runs yield + - ``gpu_to_gpu`` + ``tie_break::prefer_smaller_index`` + ``stable_sorted``: both runs yield ``[10@0, 8@1, 8@2]`` * - Tie-break changes the set, not just the order - Compare ``prefer_smaller_index`` vs ``prefer_larger_index``: ``8@2`` vs ``8@3`` @@ -273,11 +311,11 @@ Choosing requirements * - Goal - Suggested ``require(...)`` * - Maximum performance, exact result unimportant - - ``determinism::not_guaranteed, output_ordering::unsorted`` + - ``determinism::not_guaranteed, tie_break::unspecified, output_ordering::unsorted`` * - Reproducible result set, order does not matter - - ``determinism::run_to_run, output_ordering::unsorted`` + - ``determinism::run_to_run, tie_break::unspecified, output_ordering::unsorted`` * - Reproducible result set with an explicit boundary policy - - ``determinism::run_to_run, tie_break::prefer_{smaller,larger}_index, output_ordering::unsorted`` + - ``determinism::gpu_to_gpu, tie_break::prefer_{smaller,larger}_index, output_ordering::unsorted`` * - Reproducible, key-sorted output - the above + ``output_ordering::sorted`` * - Reproducible, key-sorted output with input-order stability among ties From d851e22b4abe6e698353ccba70513c2ba70360ab Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 23 Jun 2026 07:05:25 -0700 Subject: [PATCH 10/18] adds documentation on function parameters --- cub/cub/device/device_batched_topk.cuh | 149 +++++++++++++++++++++++++ 1 file changed, 149 insertions(+) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 0b379d410501..460b22d6e5a8 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -184,6 +184,9 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! DeviceBatchedTopK provides device-wide, parallel operations for finding the largest (or smallest) K items from //! many segments of unordered data items residing within device-accessible memory. //! +//! .. versionadded:: 3.5.0 +//! First appears in CUDA Toolkit 13.5. +//! //! Overview //! ++++++++++++++++++++++++++ //! @@ -404,6 +407,36 @@ struct DeviceBatchedTopK //! :end-before: example-end batched-topk-max-keys-env //! //! @endrst + //! + //! @tparam KeyInputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator + //! + //! @tparam KeyOutputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator + //! + //! @param[in] d_keys_in + //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` + //! + //! @param[out] d_keys_out + //! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i` + //! + //! @param[in] segment_sizes + //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! *Choosing argument bounds* section). + //! + //! @param[in] k + //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! + //! @param[in] num_segments + //! Annotated argument providing the (uniform) number of segments + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, + //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! @endrst template ` for a uniform size, + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! *Choosing argument bounds* section). + //! + //! @param[in] k + //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! + //! @param[in] num_segments + //! Annotated argument providing the (uniform) number of segments + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, + //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! @endrst template ` for a uniform size, + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! *Choosing argument bounds* section). + //! + //! @param[in] k + //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! + //! @param[in] num_segments + //! Annotated argument providing the (uniform) number of segments + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, + //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! @endrst template ` for a uniform size, + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! *Choosing argument bounds* section). + //! + //! @param[in] k + //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! + //! @param[in] num_segments + //! Annotated argument providing the (uniform) number of segments + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, + //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! @endrst template Date: Tue, 23 Jun 2026 07:05:44 -0700 Subject: [PATCH 11/18] validates unwrapped parameters --- cub/cub/device/device_batched_topk.cuh | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 460b22d6e5a8..bdfb635d4c63 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -31,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -154,6 +155,20 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( "cub::DeviceBatchedTopK currently requires a single (uniform) number of segments resolved on the " "host; pass num_segments as a single-value annotation (e.g. cuda::args::constant or " "cuda::args::immediate), not a per-segment sequence."); + static_assert(::cuda::args::__is_wrapper_v + || ::cuda::std::is_integral_v, + "cub::DeviceBatchedTopK: segment_sizes must be a cuda::args annotation or a plain integral value " + "(taken as a uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap " + "per-segment sizes in cuda::args::deferred_sequence, or a single device-side value in " + "cuda::args::deferred."); + static_assert(::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, + "cub::DeviceBatchedTopK: k must be a cuda::args annotation or a plain integral value (taken as a " + "uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap a per-segment k " + "in cuda::args::deferred_sequence, or a single device-side value in cuda::args::deferred."); + static_assert(::cuda::args::__is_wrapper_v + || ::cuda::std::is_integral_v, + "cub::DeviceBatchedTopK: num_segments must be a cuda::args annotation or a plain integral value. A " + "raw pointer or iterator is not accepted."); const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); @@ -211,6 +226,11 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! example one produced on the device by a preceding launch. //! - ``cuda::args::deferred_sequence{iterator}`` for a distinct value per segment, also read in stream order. //! +//! A plain integral value works too and is taken as a uniform ``immediate`` (no extra bounds). A pointer or iterator, +//! by contrast, must be wrapped explicitly in ``deferred`` (single value) or ``deferred_sequence`` (per segment). +//! Passing a raw pointer or iterator is rejected at compile time, because it would otherwise be misread as a single +//! value rather than a sequence. +//! //! **How it is bounded.** A bound lets the algorithm reason about a value it does not know exactly: //! //! - A **compile-time** bound, ``cuda::args::bounds()``, may accompany ``immediate``, ``deferred``, or From 6cb0864cd5696fe53f85fb26c551e4ab91ff6d07 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 23 Jun 2026 09:15:31 -0700 Subject: [PATCH 12/18] clarifies support for plain integrals --- cub/cub/device/device_batched_topk.cuh | 29 +++++++++++++------------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index bdfb635d4c63..c6e977c79feb 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -240,10 +240,11 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! the range is only known at runtime. When combined with a compile-time bound, the runtime bound must be at least //! as narrow, lying within the compile-time range and only tightening it further. //! -//! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. The kernel specializes -//! on their upper bound, so each carries a compile-time bound (a ``constant`` is its own, the other forms take an -//! explicit ``cuda::args::bounds()``). ``num_segments`` must be a single value known on the host, supplied -//! as ``constant`` or ``immediate``. +//! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. Only ``segment_sizes`` +//! needs a small compile-time upper bound that the kernel specializes on, so it must be a ``constant`` or carry an +//! explicit ``cuda::args::bounds()`` (a plain integral value is not enough). ``k`` needs no bound, since it is +//! capped per segment to the segment size, so a plain integral value or any annotation works. ``num_segments`` is a +//! single host value given as a ``constant``, an ``immediate``, or a plain integral, and needs no bound. //! //! .. code-block:: c++ //! @@ -367,10 +368,10 @@ struct DeviceBatchedTopK //! *Choosing argument bounds* section). //! //! @param[in] k - //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! Annotated argument providing the (uniform) number of segments + //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] env //! @rst @@ -447,10 +448,10 @@ struct DeviceBatchedTopK //! *Choosing argument bounds* section). //! //! @param[in] k - //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! Annotated argument providing the (uniform) number of segments + //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] env //! @rst @@ -577,10 +578,10 @@ struct DeviceBatchedTopK //! *Choosing argument bounds* section). //! //! @param[in] k - //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! Annotated argument providing the (uniform) number of segments + //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] env //! @rst @@ -737,10 +738,10 @@ struct DeviceBatchedTopK //! *Choosing argument bounds* section). //! //! @param[in] k - //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! Annotated argument providing the (uniform) number of segments + //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] env //! @rst @@ -876,10 +877,10 @@ struct DeviceBatchedTopK //! *Choosing argument bounds* section). //! //! @param[in] k - //! Annotated argument providing the number of selected items per segment. Capped per segment to the segment size. + //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] num_segments - //! Annotated argument providing the (uniform) number of segments + //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. //! //! @param[in] env //! @rst From b9ffd437d1bfc7f8e03ca6893ea92eb7978d0327 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 23 Jun 2026 09:15:39 -0700 Subject: [PATCH 13/18] fixes copydetails --- cub/cub/device/device_batched_topk.cuh | 249 +++++++++++++++++++++++-- 1 file changed, 231 insertions(+), 18 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index c6e977c79feb..a37a898df53b 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -155,20 +155,20 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( "cub::DeviceBatchedTopK currently requires a single (uniform) number of segments resolved on the " "host; pass num_segments as a single-value annotation (e.g. cuda::args::constant or " "cuda::args::immediate), not a per-segment sequence."); - static_assert(::cuda::args::__is_wrapper_v - || ::cuda::std::is_integral_v, - "cub::DeviceBatchedTopK: segment_sizes must be a cuda::args annotation or a plain integral value " - "(taken as a uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap " - "per-segment sizes in cuda::args::deferred_sequence, or a single device-side value in " - "cuda::args::deferred."); + static_assert( + ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, + "cub::DeviceBatchedTopK: segment_sizes must be a cuda::args annotation or a plain integral value " + "(taken as a uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap " + "per-segment sizes in cuda::args::deferred_sequence, or a single device-side value in " + "cuda::args::deferred."); static_assert(::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, "cub::DeviceBatchedTopK: k must be a cuda::args annotation or a plain integral value (taken as a " "uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap a per-segment k " "in cuda::args::deferred_sequence, or a single device-side value in cuda::args::deferred."); - static_assert(::cuda::args::__is_wrapper_v - || ::cuda::std::is_integral_v, - "cub::DeviceBatchedTopK: num_segments must be a cuda::args annotation or a plain integral value. A " - "raw pointer or iterator is not accepted."); + static_assert( + ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, + "cub::DeviceBatchedTopK: num_segments must be a cuda::args annotation or a plain integral value. A " + "raw pointer or iterator is not accepted."); const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); @@ -240,11 +240,10 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( //! the range is only known at runtime. When combined with a compile-time bound, the runtime bound must be at least //! as narrow, lying within the compile-time range and only tightening it further. //! -//! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. Only ``segment_sizes`` -//! needs a small compile-time upper bound that the kernel specializes on, so it must be a ``constant`` or carry an -//! explicit ``cuda::args::bounds()`` (a plain integral value is not enough). ``k`` needs no bound, since it is -//! capped per segment to the segment size, so a plain integral value or any annotation works. ``num_segments`` is a -//! single host value given as a ``constant``, an ``immediate``, or a plain integral, and needs no bound. +//! **Which form each parameter accepts.** ``segment_sizes`` and ``k`` accept all four forms. ``num_segments`` must be +//! a single value (``constant``, ``immediate``, or a plain integral), never a per-segment sequence. ``segment_sizes`` +//! must also carry a small compile-time upper bound (a ``constant`` or ``cuda::args::bounds()``), and tight +//! bounds on every parameter are encouraged. //! //! .. code-block:: c++ //! @@ -348,6 +347,18 @@ struct DeviceBatchedTopK //! @tparam KeyOutputIteratorItT //! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator //! + //! @tparam SegmentSizeParameterT + //! **[inferred]** Type of the ``segment_sizes`` argument + //! + //! @tparam KParameterT + //! **[inferred]** Type of the ``k`` argument + //! + //! @tparam NumSegmentsParameterT + //! **[inferred]** Type of the ``num_segments`` argument + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! //! @param[in] d_temp_storage //! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to //! `temp_storage_bytes` and no work is done. @@ -435,6 +446,18 @@ struct DeviceBatchedTopK //! @tparam KeyOutputIteratorItT //! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator //! + //! @tparam SegmentSizeParameterT + //! **[inferred]** Type of the ``segment_sizes`` argument + //! + //! @tparam KParameterT + //! **[inferred]** Type of the ``k`` argument + //! + //! @tparam NumSegmentsParameterT + //! **[inferred]** Type of the ``num_segments`` argument + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! //! @param[in] d_keys_in //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` //! @@ -509,7 +532,54 @@ struct DeviceBatchedTopK //! //! @endrst //! - //! @copydetails MaxKeys + //! @tparam KeyInputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator + //! + //! @tparam KeyOutputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator + //! + //! @tparam SegmentSizeParameterT + //! **[inferred]** Type of the ``segment_sizes`` argument + //! + //! @tparam KParameterT + //! **[inferred]** Type of the ``k`` argument + //! + //! @tparam NumSegmentsParameterT + //! **[inferred]** Type of the ``num_segments`` argument + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_temp_storage + //! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to + //! `temp_storage_bytes` and no work is done. + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] d_keys_in + //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` + //! + //! @param[out] d_keys_out + //! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i` + //! + //! @param[in] segment_sizes + //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! *Choosing argument bounds* section). + //! + //! @param[in] k + //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. + //! + //! @param[in] num_segments + //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, + //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! @endrst template ``. + //! //! @param[in] d_keys_in //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` //! @@ -643,12 +725,43 @@ struct DeviceBatchedTopK //! //! @endrst //! + //! @tparam KeyInputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator + //! + //! @tparam KeyOutputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator + //! //! @tparam ValueInputIteratorItT //! **[inferred]** Random-access input iterator over per-segment value-input iterators @iterator //! //! @tparam ValueOutputIteratorItT //! **[inferred]** Random-access input iterator over per-segment value-output iterators @iterator //! + //! @tparam SegmentSizeParameterT + //! **[inferred]** Type of the ``segment_sizes`` argument + //! + //! @tparam KParameterT + //! **[inferred]** Type of the ``k`` argument + //! + //! @tparam NumSegmentsParameterT + //! **[inferred]** Type of the ``num_segments`` argument + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_temp_storage + //! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to + //! `temp_storage_bytes` and no work is done. + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] d_keys_in + //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` + //! + //! @param[out] d_keys_out + //! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i` + //! //! @param[in] d_values_in //! Iterator such that `d_values_in[i]` yields a random-access iterator to the values of segment `i` //! @@ -656,7 +769,23 @@ struct DeviceBatchedTopK //! Iterator such that `d_values_out[i]` yields a random-access output iterator for the values corresponding to the //! top-k keys of segment `i` //! - //! @copydetails MaxKeys + //! @param[in] segment_sizes + //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! *Choosing argument bounds* section). + //! + //! @param[in] k + //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. + //! + //! @param[in] num_segments + //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, + //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! @endrst template ``. + //! //! @param[in] d_keys_in //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` //! @@ -795,7 +936,67 @@ struct DeviceBatchedTopK //! //! @endrst //! - //! @copydetails MaxPairs + //! @tparam KeyInputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator + //! + //! @tparam KeyOutputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator + //! + //! @tparam ValueInputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment value-input iterators @iterator + //! + //! @tparam ValueOutputIteratorItT + //! **[inferred]** Random-access input iterator over per-segment value-output iterators @iterator + //! + //! @tparam SegmentSizeParameterT + //! **[inferred]** Type of the ``segment_sizes`` argument + //! + //! @tparam KParameterT + //! **[inferred]** Type of the ``k`` argument + //! + //! @tparam NumSegmentsParameterT + //! **[inferred]** Type of the ``num_segments`` argument + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_temp_storage + //! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to + //! `temp_storage_bytes` and no work is done. + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] d_keys_in + //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` + //! + //! @param[out] d_keys_out + //! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i` + //! + //! @param[in] d_values_in + //! Iterator such that `d_values_in[i]` yields a random-access iterator to the values of segment `i` + //! + //! @param[out] d_values_out + //! Iterator such that `d_values_out[i]` yields a random-access output iterator for the values corresponding to the + //! top-k keys of segment `i` + //! + //! @param[in] segment_sizes + //! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant` for a uniform size, + //! or `cuda::args::deferred_sequence{...}` for variable sizes). Must carry a small compile-time maximum. + //! Prefer a sharp (tight) upper bound, since a looser bound may increase temporary-storage usage (see the + //! *Choosing argument bounds* section). + //! + //! @param[in] k + //! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value. + //! + //! @param[in] num_segments + //! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Must require `determinism::not_guaranteed`, + //! `tie_break::unspecified`, and `output_ordering::unsorted`. + //! @endrst template ``. + //! //! @param[in] d_keys_in //! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i` //! From 4cbb92980e356688a0fe356fce12f69675136bae Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Sat, 27 Jun 2026 08:02:36 -0700 Subject: [PATCH 14/18] adds api test for unwrapped integrals --- .../catch2_test_device_segmented_topk_keys.cu | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 86c1a00a0267..0c08d9811206 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -508,3 +508,53 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys preserve -0.0f in output", const int num_minus_zero = static_cast(thrust::count_if(d_keys_out.begin(), d_keys_out.end(), is_minus_zero{})); REQUIRE(num_minus_zero >= 1); } + +// Users may pass `k` and `num_segments` un-annotated. A plain integral value is taken as a uniform immediate with no +// compile-time bound. +C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys accept unwrapped (plain integral) k and num_segments", + "[keys][segmented][topk][device]", + key_types, + max_segment_size_list, + max_num_k_list, + select_direction_list) +{ + using segment_size_t = cuda::std::int64_t; + using segment_index_t = cuda::std::int64_t; + + using key_t = c2h::get<0, TestType>; + + constexpr segment_size_t static_max_segment_size = c2h::get<1, TestType>::value; + constexpr segment_size_t static_max_k = c2h::get<2, TestType>::value; + constexpr auto direction = c2h::get<3, TestType>::value; + + // Fixed sizes: this test exercises the argument form, not the size matrix. + const segment_size_t segment_size = (cuda::std::min) (segment_size_t{256}, static_max_segment_size); + const segment_size_t k = (cuda::std::min) (static_max_k, segment_size); + const segment_index_t num_segments = 42; + + CAPTURE(c2h::type_name(), static_max_segment_size, static_max_k, segment_size, k, num_segments, direction); + + c2h::device_vector keys_in_buffer(num_segments * segment_size, thrust::no_init); + c2h::device_vector keys_out_buffer(num_segments * k, thrust::no_init); + c2h::gen(C2H_SEED(1), keys_in_buffer); + auto d_keys_in_ptr = thrust::raw_pointer_cast(keys_in_buffer.data()); + auto d_keys_out_ptr = thrust::raw_pointer_cast(keys_out_buffer.data()); + auto d_keys_in = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_in_ptr), segment_size); + auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(d_keys_out_ptr), k); + + c2h::device_vector expected_keys(keys_in_buffer); + + // `k` and `num_segments` are passed as plain integral values (un-annotated immediate). + batched_topk_keys( + d_keys_in, + d_keys_out, + cuda::args::immediate{segment_size, cuda::args::bounds()}, + k, + num_segments); + + fixed_size_segmented_sort_keys(expected_keys, num_segments, segment_size, direction); + compact_sorted_keys_to_topk(expected_keys, segment_size, k); + fixed_size_segmented_sort_keys(keys_out_buffer, num_segments, k, direction); + + REQUIRE(expected_keys == keys_out_buffer); +} From 6b0e679f960698425ed4ea0242ce28349e2dff77 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Sat, 27 Jun 2026 08:38:08 -0700 Subject: [PATCH 15/18] addresses review comments --- cub/cub/device/device_batched_topk.cuh | 74 +++++---- .../catch2_test_device_batched_topk_api.cu | 140 +++++++----------- ...catch2_test_device_batched_topk_env_api.cu | 18 --- .../catch2_test_device_segmented_topk_keys.cu | 2 +- 4 files changed, 91 insertions(+), 143 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index a37a898df53b..6366228fe3ca 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -175,7 +175,7 @@ CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( // The total-number-of-items guarantee is intentionally not part of the initial public API surface. The dispatch // only uses its element type to size internal large-segment offsets (the value itself is unused), so we pass a // conservative 64-bit upper bound here. - const auto total_num_items = ::cuda::args::immediate{::cuda::std::numeric_limits<::cuda::std::int64_t>::max()}; + constexpr auto total_num_items = ::cuda::args::immediate{::cuda::std::numeric_limits<::cuda::std::int64_t>::max()}; return batched_topk::dispatch( d_temp_storage, @@ -496,20 +496,19 @@ struct DeviceBatchedTopK EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxKeys"); - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { - return detail::dispatch_batched_topk( - storage, - bytes, - d_keys_in, - d_keys_out, - static_cast(nullptr), - static_cast(nullptr), - segment_sizes, - k, - num_segments, - env); - }); + return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) { + return detail::dispatch_batched_topk( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + num_segments, + env); + }); } //! @rst @@ -685,20 +684,19 @@ struct DeviceBatchedTopK EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinKeys"); - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { - return detail::dispatch_batched_topk( - storage, - bytes, - d_keys_in, - d_keys_out, - static_cast(nullptr), - static_cast(nullptr), - segment_sizes, - k, - num_segments, - env); - }); + return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) { + return detail::dispatch_batched_topk( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + segment_sizes, + k, + num_segments, + env); + }); } //! @rst @@ -908,11 +906,10 @@ struct DeviceBatchedTopK EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxPairs"); - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { - return detail::dispatch_batched_topk( - storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env); - }); + return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) { + return detail::dispatch_batched_topk( + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env); + }); } //! @rst @@ -1119,11 +1116,10 @@ struct DeviceBatchedTopK EnvT env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinPairs"); - return detail::dispatch_with_env( - env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) { - return detail::dispatch_batched_topk( - storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env); - }); + return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) { + return detail::dispatch_batched_topk( + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env); + }); } }; diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu index 9061f55001b4..aa8fd03fa7c8 100644 --- a/cub/test/catch2_test_device_batched_topk_api.cu +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -18,8 +18,6 @@ #include #include -#include - #include C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_topk][device]") @@ -40,9 +38,9 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_t // Argument annotations: a small, compile-time segment size and k, plus the runtime segment count and item-count // bound. - auto segment_sizes = cuda::args::constant{}; - auto k_arg = cuda::args::constant{}; - auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + constexpr auto segment_sizes = cuda::args::constant{}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; // Top-k output is unordered and may be non-deterministic; this must be acknowledged via the environment. auto env = cuda::std::execution::env{cuda::execution::require( @@ -57,22 +55,15 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_t // Allocate temporary storage and run thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - if (error == cudaSuccess) - { - error = cub::DeviceBatchedTopK::MaxKeys( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - segment_sizes, - k_arg, - num_segs, - env); - } - if (error != cudaSuccess) - { - std::cerr << "cub::DeviceBatchedTopK::MaxKeys failed with status: " << error << '\n'; - } + error = cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); // example-end batched-topk-max-keys REQUIRE(error == cudaSuccess); @@ -97,10 +88,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys temp-storage API example", "[batched_t auto d_keys_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); - auto segment_sizes = cuda::args::constant{}; - auto k_arg = cuda::args::constant{}; - auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto env = cuda::std::execution::env{cuda::execution::require( + constexpr auto segment_sizes = cuda::args::constant{}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto env = cuda::std::execution::env{cuda::execution::require( cuda::execution::determinism::not_guaranteed, cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)}; @@ -109,22 +100,15 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys temp-storage API example", "[batched_t auto error = cub::DeviceBatchedTopK::MinKeys( nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - if (error == cudaSuccess) - { - error = cub::DeviceBatchedTopK::MinKeys( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - segment_sizes, - k_arg, - num_segs, - env); - } - if (error != cudaSuccess) - { - std::cerr << "cub::DeviceBatchedTopK::MinKeys failed with status: " << error << '\n'; - } + error = cub::DeviceBatchedTopK::MinKeys( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); // example-end batched-topk-min-keys REQUIRE(error == cudaSuccess); @@ -153,10 +137,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); - auto segment_sizes = cuda::args::constant{}; - auto k_arg = cuda::args::constant{}; - auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto env = cuda::std::execution::env{cuda::execution::require( + constexpr auto segment_sizes = cuda::args::constant{}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto env = cuda::std::execution::env{cuda::execution::require( cuda::execution::determinism::not_guaranteed, cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)}; @@ -165,24 +149,17 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ auto error = cub::DeviceBatchedTopK::MaxPairs( nullptr, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k_arg, num_segs, env); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - if (error == cudaSuccess) - { - error = cub::DeviceBatchedTopK::MaxPairs( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k_arg, - num_segs, - env); - } - if (error != cudaSuccess) - { - std::cerr << "cub::DeviceBatchedTopK::MaxPairs failed with status: " << error << '\n'; - } + error = cub::DeviceBatchedTopK::MaxPairs( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs, + env); // example-end batched-topk-max-pairs REQUIRE(error == cudaSuccess); @@ -229,10 +206,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ auto d_values_out = cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(values_out.data())), k); - auto segment_sizes = cuda::args::constant{}; - auto k_arg = cuda::args::constant{}; - auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; - auto env = cuda::std::execution::env{cuda::execution::require( + constexpr auto segment_sizes = cuda::args::constant{}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto env = cuda::std::execution::env{cuda::execution::require( cuda::execution::determinism::not_guaranteed, cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)}; @@ -241,24 +218,17 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ auto error = cub::DeviceBatchedTopK::MinPairs( nullptr, temp_storage_bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k_arg, num_segs, env); thrust::device_vector temp_storage(temp_storage_bytes, thrust::no_init); - if (error == cudaSuccess) - { - error = cub::DeviceBatchedTopK::MinPairs( - thrust::raw_pointer_cast(temp_storage.data()), - temp_storage_bytes, - d_keys_in, - d_keys_out, - d_values_in, - d_values_out, - segment_sizes, - k_arg, - num_segs, - env); - } - if (error != cudaSuccess) - { - std::cerr << "cub::DeviceBatchedTopK::MinPairs failed with status: " << error << '\n'; - } + error = cub::DeviceBatchedTopK::MinPairs( + thrust::raw_pointer_cast(temp_storage.data()), + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + segment_sizes, + k_arg, + num_segs, + env); // example-end batched-topk-min-pairs REQUIRE(error == cudaSuccess); diff --git a/cub/test/catch2_test_device_batched_topk_env_api.cu b/cub/test/catch2_test_device_batched_topk_env_api.cu index 0fa51808e166..694a0093668f 100644 --- a/cub/test/catch2_test_device_batched_topk_env_api.cu +++ b/cub/test/catch2_test_device_batched_topk_env_api.cu @@ -20,8 +20,6 @@ #include #include -#include - #include C2H_TEST("cub::DeviceBatchedTopK::MaxKeys env-alloc example", "[batched_topk][device][env]") @@ -54,10 +52,6 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys env-alloc example", "[batched_topk][de cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, env); - if (error != cudaSuccess) - { - std::cerr << "cub::DeviceBatchedTopK::MaxKeys failed with status: " << error << '\n'; - } // example-end batched-topk-max-keys-env stream.sync(); @@ -96,10 +90,6 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys env-alloc example", "[batched_topk][de cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, env); - if (error != cudaSuccess) - { - std::cerr << "cub::DeviceBatchedTopK::MinKeys failed with status: " << error << '\n'; - } // example-end batched-topk-min-keys-env stream.sync(); @@ -144,10 +134,6 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][d cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, env); - if (error != cudaSuccess) - { - std::cerr << "cub::DeviceBatchedTopK::MaxPairs failed with status: " << error << '\n'; - } // example-end batched-topk-max-pairs-env stream.sync(); @@ -211,10 +197,6 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][d cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, env); - if (error != cudaSuccess) - { - std::cerr << "cub::DeviceBatchedTopK::MinPairs failed with status: " << error << '\n'; - } // example-end batched-topk-min-pairs-env stream.sync(); diff --git a/cub/test/catch2_test_device_segmented_topk_keys.cu b/cub/test/catch2_test_device_segmented_topk_keys.cu index 0c08d9811206..8e1d6ba29f3e 100644 --- a/cub/test/catch2_test_device_segmented_topk_keys.cu +++ b/cub/test/catch2_test_device_segmented_topk_keys.cu @@ -510,7 +510,7 @@ C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys preserve -0.0f in output", } // Users may pass `k` and `num_segments` un-annotated. A plain integral value is taken as a uniform immediate with no -// compile-time bound. +// compile-time bound. C2H_TEST("DeviceBatchedTopK::{Min,Max}Keys accept unwrapped (plain integral) k and num_segments", "[keys][segmented][topk][device]", key_types, From 840ebdde951aba3e0cd317017af1cd162d3baef2 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 30 Jun 2026 00:19:55 -0700 Subject: [PATCH 16/18] adds to the umbrella cub.cuh and device-wide docs --- cub/cub/cub.cuh | 1 + docs/cub/api_docs/device_wide.rst | 1 + 2 files changed, 2 insertions(+) diff --git a/cub/cub/cub.cuh b/cub/cub/cub.cuh index 66d7ca6434c3..9d1abbc7b8fd 100644 --- a/cub/cub/cub.cuh +++ b/cub/cub/cub.cuh @@ -43,6 +43,7 @@ // Device #include +#include #include #include #include diff --git a/docs/cub/api_docs/device_wide.rst b/docs/cub/api_docs/device_wide.rst index e473cef8a665..2711244181ad 100644 --- a/docs/cub/api_docs/device_wide.rst +++ b/docs/cub/api_docs/device_wide.rst @@ -110,6 +110,7 @@ CUB device-level segmented-problem (batched) parallel algorithms: * :cpp:struct:`cub::DeviceSegmentedSort` computes batched sort across non-overlapping sequences of data residing within device-accessible memory * :cpp:struct:`cub::DeviceSegmentedRadixSort` computes batched radix sort across non-overlapping sequences of data residing within device-accessible memory * :cpp:struct:`cub::DeviceSegmentedReduce` computes reductions across multiple sequences of data residing within device-accessible memory +* :cpp:struct:`cub::DeviceBatchedTopK` finds the largest (or smallest) K items from each of multiple unordered lists (segments) residing within device-accessible memory * :cpp:struct:`cub::DeviceCopy` provides device-wide, parallel operations for batched copying of data residing within device-accessible memory * :cpp:struct:`cub::DeviceMemcpy` provides device-wide, parallel operations for batched copying of data residing within device-accessible memory * :cpp:struct:`cub::DeviceFind` provides vectorized binary search algorithms From 484c418d262e089e844129a149e0f572bbe11a46 Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 30 Jun 2026 00:21:19 -0700 Subject: [PATCH 17/18] const-rev env and docs on in/out overlapping --- cub/cub/device/device_batched_topk.cuh | 34 +++++++++++++++----------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/cub/cub/device/device_batched_topk.cuh b/cub/cub/device/device_batched_topk.cuh index 6366228fe3ca..54261ac9590b 100644 --- a/cub/cub/device/device_batched_topk.cuh +++ b/cub/cub/device/device_batched_topk.cuh @@ -326,7 +326,8 @@ struct DeviceBatchedTopK //! //! .. note:: //! - //! The behavior is undefined if the input and output ranges overlap in any way. + //! The behavior is undefined if an output range overlaps another output range or any input range. + //! Input ranges may overlap one another. //! //! - @devicestorage //! @@ -403,7 +404,7 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - EnvT env = {}) + const EnvT& env = {}) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MaxKeys"); return detail::dispatch_batched_topk( @@ -427,7 +428,8 @@ struct DeviceBatchedTopK //! //! .. note:: //! - //! The behavior is undefined if the input and output ranges overlap in any way. + //! The behavior is undefined if an output range overlaps another output range or any input range. + //! Input ranges may overlap one another. //! //! Snippet //! +++++++++++++++++++++++++++++++++++++++++++++ @@ -493,7 +495,7 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - EnvT env = {}) + const EnvT& env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxKeys"); return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) { @@ -516,7 +518,8 @@ struct DeviceBatchedTopK //! //! .. note:: //! - //! The behavior is undefined if the input and output ranges overlap in any way. + //! The behavior is undefined if an output range overlaps another output range or any input range. + //! Input ranges may overlap one another. //! //! - @devicestorage //! @@ -593,7 +596,7 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - EnvT env = {}) + const EnvT& env = {}) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MinKeys"); return detail::dispatch_batched_topk( @@ -615,7 +618,8 @@ struct DeviceBatchedTopK //! //! .. note:: //! - //! The behavior is undefined if the input and output ranges overlap in any way. + //! The behavior is undefined if an output range overlaps another output range or any input range. + //! Input ranges may overlap one another. //! //! Snippet //! +++++++++++++++++++++++++++++++++++++++++++++ @@ -681,7 +685,7 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - EnvT env = {}) + const EnvT& env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinKeys"); return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) { @@ -708,7 +712,8 @@ struct DeviceBatchedTopK //! //! .. note:: //! - //! The behavior is undefined if the input and output ranges overlap in any way. + //! The behavior is undefined if an output range overlaps another output range or any input range. + //! Input ranges may overlap one another. //! //! - @devicestorage //! @@ -802,7 +807,7 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - EnvT env = {}) + const EnvT& env = {}) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MaxPairs"); return detail::dispatch_batched_topk( @@ -903,7 +908,7 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - EnvT env = {}) + const EnvT& env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxPairs"); return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) { @@ -918,7 +923,8 @@ struct DeviceBatchedTopK //! //! .. note:: //! - //! The behavior is undefined if the input and output ranges overlap in any way. + //! The behavior is undefined if an output range overlaps another output range or any input range. + //! Input ranges may overlap one another. //! //! - @devicestorage //! @@ -1012,7 +1018,7 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - EnvT env = {}) + const EnvT& env = {}) { _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MinPairs"); return detail::dispatch_batched_topk( @@ -1113,7 +1119,7 @@ struct DeviceBatchedTopK SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, - EnvT env = {}) + const EnvT& env = {}) { _CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinPairs"); return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) { From 4c7a6e137358f362f4faa39bb05644680d0ef9df Mon Sep 17 00:00:00 2001 From: Elias Stehle <3958403+elstehle@users.noreply.github.com> Date: Tue, 30 Jun 2026 00:42:28 -0700 Subject: [PATCH 18/18] improves api examples; adds mis-aligned temp storage test --- .../catch2_test_device_batched_topk_api.cu | 82 ++++++++++++++++--- ...catch2_test_device_batched_topk_env_api.cu | 34 +++++--- 2 files changed, 95 insertions(+), 21 deletions(-) diff --git a/cub/test/catch2_test_device_batched_topk_api.cu b/cub/test/catch2_test_device_batched_topk_api.cu index aa8fd03fa7c8..92d55da4adea 100644 --- a/cub/test/catch2_test_device_batched_topk_api.cu +++ b/cub/test/catch2_test_device_batched_topk_api.cu @@ -64,13 +64,16 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys temp-storage API example", "[batched_t k_arg, num_segs, env); + // Each segment's k largest keys are written to keys_out in unspecified order. The result set is fixed, + // shown here sorted per segment: + auto expected_result_set = thrust::device_vector{8, 7, 6, /* segment 0 */ 9, 8, 7 /* segment 1 */}; // example-end batched-topk-max-keys REQUIRE(error == cudaSuccess); - // The per-segment output is unordered; sort each segment (descending) before comparing. + // keys_out is unordered, so sort each segment (descending) before comparing against the expected set. thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); - REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); + REQUIRE(keys_out == expected_result_set); } C2H_TEST("cub::DeviceBatchedTopK::MinKeys temp-storage API example", "[batched_topk][device]") @@ -109,12 +112,16 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys temp-storage API example", "[batched_t k_arg, num_segs, env); + // Each segment's k smallest keys are written to keys_out in unspecified order. The result set is fixed, + // shown here sorted per segment: + auto expected_result_set = thrust::device_vector{-3, 1, 2, /* segment 0 */ 0, 1, 2 /* segment 1 */}; // example-end batched-topk-min-keys REQUIRE(error == cudaSuccess); + // keys_out is unordered, so sort each segment (ascending) before comparing against the expected set. thrust::sort(keys_out.begin(), keys_out.begin() + k); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); - REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); + REQUIRE(keys_out == expected_result_set); } C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_topk][device]") @@ -160,12 +167,15 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ k_arg, num_segs, env); + // keys_out holds each segment's k largest keys. The key set is fixed (shown here sorted per segment). For + // keys that tie, which equal element's value is returned is unspecified. + auto expected_result_set = thrust::device_vector{8, 7, 6, /* segment 0 */ 9, 8, 7 /* segment 1 */}; // example-end batched-topk-max-pairs REQUIRE(error == cudaSuccess); - // Each input value is the global item index. Bounds-checked: every returned value indexes back to the input - // element (within its own segment) whose key was selected. + // Each returned value is the source index of its key within the segment. Check that every value indexes + // back to the input element whose key was selected. thrust::host_vector h_keys_in(keys_in); thrust::host_vector h_keys_out(keys_out); thrust::host_vector h_values_out(values_out); @@ -181,10 +191,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs temp-storage API example", "[batched_ } } - // The selected keys must be the per-segment top-k (output is unordered; sort each segment descending). + // keys_out is unordered, so sort each segment (descending) before comparing against the expected set. thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); - REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); + REQUIRE(keys_out == expected_result_set); } C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_topk][device]") @@ -229,12 +239,15 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ k_arg, num_segs, env); + // keys_out holds each segment's k smallest keys. The key set is fixed (shown here sorted per segment). For + // keys that tie, which equal element's value is returned is unspecified. + auto expected_result_set = thrust::device_vector{-3, 1, 2, /* segment 0 */ 0, 1, 2 /* segment 1 */}; // example-end batched-topk-min-pairs REQUIRE(error == cudaSuccess); - // Each input value is the global item index. Bounds-checked: every returned value indexes back to the input - // element (within its own segment) whose key was selected. + // Each returned value is the source index of its key within the segment. Check that every value indexes + // back to the input element whose key was selected. thrust::host_vector h_keys_in(keys_in); thrust::host_vector h_keys_out(keys_out); thrust::host_vector h_values_out(values_out); @@ -250,8 +263,55 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs temp-storage API example", "[batched_ } } - // The selected keys must be the per-segment top-k (output is unordered; sort each segment ascending). + // keys_out is unordered, so sort each segment (ascending) before comparing against the expected set. thrust::sort(keys_out.begin(), keys_out.begin() + k); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); - REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); + REQUIRE(keys_out == expected_result_set); +} + +// The temporary storage size requirement must not assume a particular base-pointer alignment (the public contract +// states that no special alignment is required). Over-allocate by one byte and offset the base pointer. +C2H_TEST("cub::DeviceBatchedTopK::MaxKeys handles a misaligned temporary storage pointer", "[batched_topk][device]") +{ + constexpr int num_segments = 2; + constexpr int segment_size = 8; + constexpr int k = 3; + + auto keys_in = thrust::device_vector{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4}; + auto keys_out = thrust::device_vector(num_segments * k, thrust::no_init); + + auto d_keys_in = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_in.data())), segment_size); + auto d_keys_out = + cuda::make_strided_iterator(cuda::make_counting_iterator(thrust::raw_pointer_cast(keys_out.data())), k); + + constexpr auto segment_sizes = cuda::args::constant{}; + constexpr auto k_arg = cuda::args::constant{}; + auto num_segs = cuda::args::immediate{cuda::std::int64_t{num_segments}}; + auto env = cuda::std::execution::env{cuda::execution::require( + cuda::execution::determinism::not_guaranteed, + cuda::execution::tie_break::unspecified, + cuda::execution::output_ordering::unsorted)}; + + size_t temp_storage_bytes = 0; + auto error = cub::DeviceBatchedTopK::MaxKeys( + nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env); + REQUIRE(error == cudaSuccess); + + // Allocate one extra byte and offset the base pointer by one to misalign it. + thrust::device_vector temp_storage(temp_storage_bytes + 1, thrust::no_init); + error = cub::DeviceBatchedTopK::MaxKeys( + thrust::raw_pointer_cast(temp_storage.data()) + 1, + temp_storage_bytes, + d_keys_in, + d_keys_out, + segment_sizes, + k_arg, + num_segs, + env); + REQUIRE(error == cudaSuccess); + + thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); + thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); + REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); } diff --git a/cub/test/catch2_test_device_batched_topk_env_api.cu b/cub/test/catch2_test_device_batched_topk_env_api.cu index 694a0093668f..f0650bde1dbb 100644 --- a/cub/test/catch2_test_device_batched_topk_env_api.cu +++ b/cub/test/catch2_test_device_batched_topk_env_api.cu @@ -52,13 +52,17 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxKeys env-alloc example", "[batched_topk][de cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, env); + // Each segment's k largest keys are written to keys_out in unspecified order. The result set is fixed, + // shown here sorted per segment: + auto expected_result_set = thrust::device_vector{8, 7, 6, /* segment 0 */ 9, 8, 7 /* segment 1 */}; // example-end batched-topk-max-keys-env stream.sync(); REQUIRE(error == cudaSuccess); + // keys_out is unordered, so sort each segment (descending) before comparing against the expected set. thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); - REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); + REQUIRE(keys_out == expected_result_set); } C2H_TEST("cub::DeviceBatchedTopK::MinKeys env-alloc example", "[batched_topk][device][env]") @@ -90,13 +94,17 @@ C2H_TEST("cub::DeviceBatchedTopK::MinKeys env-alloc example", "[batched_topk][de cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, env); + // Each segment's k smallest keys are written to keys_out in unspecified order. The result set is fixed, + // shown here sorted per segment: + auto expected_result_set = thrust::device_vector{-3, 1, 2, /* segment 0 */ 0, 1, 2 /* segment 1 */}; // example-end batched-topk-min-keys-env stream.sync(); REQUIRE(error == cudaSuccess); + // keys_out is unordered, so sort each segment (ascending) before comparing against the expected set. thrust::sort(keys_out.begin(), keys_out.begin() + k); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); - REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); + REQUIRE(keys_out == expected_result_set); } C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][device][env]") @@ -134,13 +142,16 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][d cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, env); + // keys_out holds each segment's k largest keys. The key set is fixed (shown here sorted per segment). For + // keys that tie, which equal element's value is returned is unspecified. + auto expected_result_set = thrust::device_vector{8, 7, 6, /* segment 0 */ 9, 8, 7 /* segment 1 */}; // example-end batched-topk-max-pairs-env stream.sync(); REQUIRE(error == cudaSuccess); - // Each input value is the global item index. Bounds-checked: every returned value indexes back to the input - // element (within its own segment) whose key was selected. + // Each returned value is the source index of its key within the segment. Check that every value indexes + // back to the input element whose key was selected. thrust::host_vector h_keys_in(keys_in); thrust::host_vector h_keys_out(keys_out); thrust::host_vector h_values_out(values_out); @@ -156,10 +167,10 @@ C2H_TEST("cub::DeviceBatchedTopK::MaxPairs env-alloc example", "[batched_topk][d } } - // The selected keys must be the per-segment top-k (output is unordered; sort each segment descending). + // keys_out is unordered, so sort each segment (descending) before comparing against the expected set. thrust::sort(keys_out.begin(), keys_out.begin() + k, cuda::std::greater{}); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k, cuda::std::greater{}); - REQUIRE(keys_out == thrust::device_vector{8, 7, 6, 9, 8, 7}); + REQUIRE(keys_out == expected_result_set); } C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][device][env]") @@ -197,13 +208,16 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][d cuda::args::constant{}, cuda::args::immediate{cuda::std::int64_t{num_segments}}, env); + // keys_out holds each segment's k smallest keys. The key set is fixed (shown here sorted per segment). For + // keys that tie, which equal element's value is returned is unspecified. + auto expected_result_set = thrust::device_vector{-3, 1, 2, /* segment 0 */ 0, 1, 2 /* segment 1 */}; // example-end batched-topk-min-pairs-env stream.sync(); REQUIRE(error == cudaSuccess); - // Each input value is the global item index. Bounds-checked: every returned value indexes back to the input - // element (within its own segment) whose key was selected. + // Each returned value is the source index of its key within the segment. Check that every value indexes + // back to the input element whose key was selected. thrust::host_vector h_keys_in(keys_in); thrust::host_vector h_keys_out(keys_out); thrust::host_vector h_values_out(values_out); @@ -219,8 +233,8 @@ C2H_TEST("cub::DeviceBatchedTopK::MinPairs env-alloc example", "[batched_topk][d } } - // The selected keys must be the per-segment top-k (output is unordered; sort each segment ascending). + // keys_out is unordered, so sort each segment (ascending) before comparing against the expected set. thrust::sort(keys_out.begin(), keys_out.begin() + k); thrust::sort(keys_out.begin() + k, keys_out.begin() + 2 * k); - REQUIRE(keys_out == thrust::device_vector{-3, 1, 2, 0, 1, 2}); + REQUIRE(keys_out == expected_result_set); }