cub::DeviceBatchedTopK#

struct DeviceBatchedTopK#

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.

Added in version 3.5.0: First appears in CUDA Toolkit 13.5.

Overview#

Given a batch of segments, DeviceBatchedTopK finds, independently for each segment, the K largest (or smallest) items.

Argument annotation framework#

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 specialize. For that reason, we encourage you to provide as much information as you have.

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<N>{} 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{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.

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<lo, hi>(), 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. 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. 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<N> or cuda::args::bounds<lo, hi>()), and tight bounds on every parameter are encouraged.

// 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 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#

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 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 Top-K: Determinism, Tie-Breaking, and Output Ordering) is not 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#

Like 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::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. 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 Top-K: Determinism, Tie-Breaking, and Output Ordering for the full requirement model, worked examples, and guidance on choosing requirements.

Usage Considerations#

Dynamic parallelism. DeviceBatchedTopK methods can be called within kernel code on devices in which CUDA dynamic parallelism is supported.

Note

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::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.

Public Static Functions

template<typename KeyInputIteratorItT, typename KeyOutputIteratorItT, typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, typename EnvT = ::cuda::std::execution::env<>> static inline 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, const EnvT &env = {} )#

Overview#

Finds, for each segment, the largest K keys from an unordered input sequence of keys.

Temporary storage for this operation. If d_temp_storage is nullptr, the required size is written to temp_storage_bytes without dereferencing iterators or launching kernels. Otherwise, d_temp_storage must point to a device-accessible allocation of at least temp_storage_bytes bytes. No special alignment is required. See Two-Phase API (explicit temporary storage management) for usage guidance.

A Simple Example#

constexpr int num_segments = 2;
constexpr int segment_size = 8;
constexpr int k            = 3;

auto keys_in  = thrust::device_vector<int>{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4};
auto keys_out = thrust::device_vector<int>(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.
constexpr auto segment_sizes = cuda::args::constant<segment_size>{};
constexpr auto k_arg         = cuda::args::constant<k>{};
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::tie_break::unspecified,
  cuda::execution::output_ordering::unsorted)};

// Query temporary storage 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_segs, env);

// Allocate temporary storage and run
thrust::device_vector<char> temp_storage(temp_storage_bytes, thrust::no_init);
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);
// 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<int>{8, 7, 6, /* segment 0 */ 9, 8, 7 /* segment 1 */};

Note

The behavior is undefined if an output range overlaps another output range or any input range. Input ranges may overlap one another.

Template Parameters:

KeyInputIteratorItT – [inferred] Random-access input iterator over per-segment key-input iterators (may be a simple pointer type)
KeyOutputIteratorItT – [inferred] Random-access input iterator over per-segment key-output iterators (may be a simple pointer type)
SegmentSizeParameterT – [inferred] Type of the segment_sizes argument
KParameterT – [inferred] Type of the k argument
NumSegmentsParameterT – [inferred] Type of the num_segments argument
EnvT – [inferred] Execution environment type. Default is cuda::std::execution::env<>.

Parameters:

d_temp_storage – [in] Device-accessible allocation of temporary storage. When nullptr, the required allocation size is written to temp_storage_bytes and no work is done.
temp_storage_bytes – [inout] Reference to size in bytes of d_temp_storage allocation
d_keys_in – [in] Iterator such that d_keys_in[i] yields a random-access iterator to the keys of segment i
d_keys_out – [out] Iterator such that d_keys_out[i] yields a random-access output iterator for the top-k keys of segment i
segment_sizes – [in] Annotated argument providing the per-segment sizes (e.g. cuda::args::constant<N> 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).
k – [in] The number of selected items per segment, given as a cuda::args annotation or a plain integral value.
num_segments – [in] The (uniform) number of segments, given as a cuda::args annotation or a plain integral value.
env – [in]
[optional] Execution environment. Must require determinism::not_guaranteed, tie_break::unspecified, and output_ordering::unsorted.

template<typename KeyInputIteratorItT, typename KeyOutputIteratorItT, typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, typename EnvT = ::cuda::std::execution::env<>> static inline cudaError_t MaxKeys( KeyInputIteratorItT d_keys_in, KeyOutputIteratorItT d_keys_out, SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, const EnvT &env = {} )#

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.

Snippet#

constexpr int num_segments = 2;
constexpr int segment_size = 8;
constexpr int k            = 3;

auto keys_in  = thrust::device_vector<int>{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4};
auto keys_out = thrust::device_vector<int>(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::tie_break::unspecified,
                           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<segment_size>{},
  cuda::args::constant<k>{},
  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<int>{8, 7, 6, /* segment 0 */ 9, 8, 7 /* segment 1 */};

Note

The behavior is undefined if an output range overlaps another output range or any input range. Input ranges may overlap one another.

Template Parameters:

KeyInputIteratorItT – [inferred] Random-access input iterator over per-segment key-input iterators (may be a simple pointer type)
KeyOutputIteratorItT – [inferred] Random-access input iterator over per-segment key-output iterators (may be a simple pointer type)
SegmentSizeParameterT – [inferred] Type of the segment_sizes argument
KParameterT – [inferred] Type of the k argument
NumSegmentsParameterT – [inferred] Type of the num_segments argument
EnvT – [inferred] Execution environment type. Default is cuda::std::execution::env<>.

Parameters:

d_keys_in – [in] Iterator such that d_keys_in[i] yields a random-access iterator to the keys of segment i
d_keys_out – [out] Iterator such that d_keys_out[i] yields a random-access output iterator for the top-k keys of segment i
segment_sizes – [in] Annotated argument providing the per-segment sizes (e.g. cuda::args::constant<N> 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).
k – [in] The number of selected items per segment, given as a cuda::args annotation or a plain integral value.
num_segments – [in] The (uniform) number of segments, given as a cuda::args annotation or a plain integral value.
env – [in]
[optional] Execution environment. Must require determinism::not_guaranteed, tie_break::unspecified, and output_ordering::unsorted.

template<typename KeyInputIteratorItT, typename KeyOutputIteratorItT, typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, typename EnvT = ::cuda::std::execution::env<>> static inline 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, const EnvT &env = {} )#

Finds, for each segment, the smallest K keys from an unordered input sequence of keys.

Temporary storage for this operation. If d_temp_storage is nullptr, the required size is written to temp_storage_bytes without dereferencing iterators or launching kernels. Otherwise, d_temp_storage must point to a device-accessible allocation of at least temp_storage_bytes bytes. No special alignment is required. See Two-Phase API (explicit temporary storage management) for usage guidance.

A Simple Example#

constexpr int num_segments = 2;
constexpr int segment_size = 8;
constexpr int k            = 3;

auto keys_in  = thrust::device_vector<int>{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4};
auto keys_out = thrust::device_vector<int>(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<segment_size>{};
constexpr auto k_arg         = cuda::args::constant<k>{};
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::MinKeys(
  nullptr, temp_storage_bytes, d_keys_in, d_keys_out, segment_sizes, k_arg, num_segs, env);
thrust::device_vector<char> temp_storage(temp_storage_bytes, thrust::no_init);
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);
// 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<int>{-3, 1, 2, /* segment 0 */ 0, 1, 2 /* segment 1 */};

Note

The behavior is undefined if an output range overlaps another output range or any input range. Input ranges may overlap one another.

Template Parameters:

KeyInputIteratorItT – [inferred] Random-access input iterator over per-segment key-input iterators (may be a simple pointer type)
KeyOutputIteratorItT – [inferred] Random-access input iterator over per-segment key-output iterators (may be a simple pointer type)
SegmentSizeParameterT – [inferred] Type of the segment_sizes argument
KParameterT – [inferred] Type of the k argument
NumSegmentsParameterT – [inferred] Type of the num_segments argument
EnvT – [inferred] Execution environment type. Default is cuda::std::execution::env<>.

Parameters:

d_temp_storage – [in] Device-accessible allocation of temporary storage. When nullptr, the required allocation size is written to temp_storage_bytes and no work is done.
temp_storage_bytes – [inout] Reference to size in bytes of d_temp_storage allocation
d_keys_in – [in] Iterator such that d_keys_in[i] yields a random-access iterator to the keys of segment i
d_keys_out – [out] Iterator such that d_keys_out[i] yields a random-access output iterator for the top-k keys of segment i
segment_sizes – [in] Annotated argument providing the per-segment sizes (e.g. cuda::args::constant<N> 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).
k – [in] The number of selected items per segment, given as a cuda::args annotation or a plain integral value.
num_segments – [in] The (uniform) number of segments, given as a cuda::args annotation or a plain integral value.
env – [in]
[optional] Execution environment. Must require determinism::not_guaranteed, tie_break::unspecified, and output_ordering::unsorted.

template<typename KeyInputIteratorItT, typename KeyOutputIteratorItT, typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, typename EnvT = ::cuda::std::execution::env<>> static inline cudaError_t MinKeys( KeyInputIteratorItT d_keys_in, KeyOutputIteratorItT d_keys_out, SegmentSizeParameterT segment_sizes, KParameterT k, NumSegmentsParameterT num_segments, const EnvT &env = {} )#

Finds, for each segment, the smallest K keys from an unordered input sequence of keys. Environment-based overload that allocates temporary storage internally.

Snippet#

constexpr int num_segments = 2;
constexpr int segment_size = 8;
constexpr int k            = 3;

auto keys_in  = thrust::device_vector<int>{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4};
auto keys_out = thrust::device_vector<int>(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::tie_break::unspecified,
                           cuda::execution::output_ordering::unsorted),
  cuda::stream_ref{stream}};

auto error = cub::DeviceBatchedTopK::MinKeys(
  d_keys_in,
  d_keys_out,
  cuda::args::constant<segment_size>{},
  cuda::args::constant<k>{},
  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<int>{-3, 1, 2, /* segment 0 */ 0, 1, 2 /* segment 1 */};

Note

The behavior is undefined if an output range overlaps another output range or any input range. Input ranges may overlap one another.

Template Parameters:

KeyInputIteratorItT – [inferred] Random-access input iterator over per-segment key-input iterators (may be a simple pointer type)
KeyOutputIteratorItT – [inferred] Random-access input iterator over per-segment key-output iterators (may be a simple pointer type)
SegmentSizeParameterT – [inferred] Type of the segment_sizes argument
KParameterT – [inferred] Type of the k argument
NumSegmentsParameterT – [inferred] Type of the num_segments argument
EnvT – [inferred] Execution environment type. Default is cuda::std::execution::env<>.

Parameters:

d_keys_in – [in] Iterator such that d_keys_in[i] yields a random-access iterator to the keys of segment i
d_keys_out – [out] Iterator such that d_keys_out[i] yields a random-access output iterator for the top-k keys of segment i
segment_sizes – [in] Annotated argument providing the per-segment sizes (e.g. cuda::args::constant<N> 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).
k – [in] The number of selected items per segment, given as a cuda::args annotation or a plain integral value.
num_segments – [in] The (uniform) number of segments, given as a cuda::args annotation or a plain integral value.
env – [in]
[optional] Execution environment. Must require determinism::not_guaranteed, tie_break::unspecified, and output_ordering::unsorted.

template<typename KeyInputIteratorItT, typename KeyOutputIteratorItT, typename ValueInputIteratorItT, typename ValueOutputIteratorItT, typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, typename EnvT = ::cuda::std::execution::env<>> static inline 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, const EnvT &env = {} )#

Overview#

Finds, for each segment, the largest K keys and their corresponding values from an unordered input sequence of key-value pairs.

Temporary storage for this operation. If d_temp_storage is nullptr, the required size is written to temp_storage_bytes without dereferencing iterators or launching kernels. Otherwise, d_temp_storage must point to a device-accessible allocation of at least temp_storage_bytes bytes. No special alignment is required. See Two-Phase API (explicit temporary storage management) for usage guidance.

A Simple Example#

constexpr int num_segments = 2;
constexpr int segment_size = 8;
constexpr int k            = 3;

auto keys_in    = thrust::device_vector<int>{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4};
auto keys_out   = thrust::device_vector<int>(num_segments * k, thrust::no_init);
auto values_out = thrust::device_vector<int>(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_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);

constexpr auto segment_sizes = cuda::args::constant<segment_size>{};
constexpr auto k_arg         = cuda::args::constant<k>{};
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::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<char> temp_storage(temp_storage_bytes, thrust::no_init);
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);
// 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<int>{8, 7, 6, /* segment 0 */ 9, 8, 7 /* segment 1 */};

Note

The behavior is undefined if an output range overlaps another output range or any input range. Input ranges may overlap one another.

Template Parameters:

KeyInputIteratorItT – [inferred] Random-access input iterator over per-segment key-input iterators (may be a simple pointer type)
KeyOutputIteratorItT – [inferred] Random-access input iterator over per-segment key-output iterators (may be a simple pointer type)
ValueInputIteratorItT – [inferred] Random-access input iterator over per-segment value-input iterators (may be a simple pointer type)
ValueOutputIteratorItT – [inferred] Random-access input iterator over per-segment value-output iterators (may be a simple pointer type)
SegmentSizeParameterT – [inferred] Type of the segment_sizes argument
KParameterT – [inferred] Type of the k argument
NumSegmentsParameterT – [inferred] Type of the num_segments argument
EnvT – [inferred] Execution environment type. Default is cuda::std::execution::env<>.

Parameters:

d_temp_storage – [in] Device-accessible allocation of temporary storage. When nullptr, the required allocation size is written to temp_storage_bytes and no work is done.
temp_storage_bytes – [inout] Reference to size in bytes of d_temp_storage allocation
d_keys_in – [in] Iterator such that d_keys_in[i] yields a random-access iterator to the keys of segment i
d_keys_out – [out] Iterator such that d_keys_out[i] yields a random-access output iterator for the top-k keys of segment i
d_values_in – [in] Iterator such that d_values_in[i] yields a random-access iterator to the values of segment i
d_values_out – [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
segment_sizes – [in] Annotated argument providing the per-segment sizes (e.g. cuda::args::constant<N> 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).
k – [in] The number of selected items per segment, given as a cuda::args annotation or a plain integral value.
num_segments – [in] The (uniform) number of segments, given as a cuda::args annotation or a plain integral value.
env – [in]
[optional] Execution environment. Must require determinism::not_guaranteed, tie_break::unspecified, and output_ordering::unsorted.

template<typename KeyInputIteratorItT, typename KeyOutputIteratorItT, typename ValueInputIteratorItT, typename ValueOutputIteratorItT, typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, typename EnvT = ::cuda::std::execution::env<>> static inline 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, const EnvT &env = {} )#

Finds, for each segment, the largest K keys and their corresponding values. Environment-based overload that allocates temporary storage internally.

Snippet#

constexpr int num_segments = 2;
constexpr int segment_size = 8;
constexpr int k            = 3;

auto keys_in    = thrust::device_vector<int>{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4};
auto keys_out   = thrust::device_vector<int>(num_segments * k, thrust::no_init);
auto values_out = thrust::device_vector<int>(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_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);

cuda::stream stream{cuda::devices[0]};
auto env = cuda::std::execution::env{
  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(
  d_keys_in,
  d_keys_out,
  d_values_in,
  d_values_out,
  cuda::args::constant<segment_size>{},
  cuda::args::constant<k>{},
  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<int>{8, 7, 6, /* segment 0 */ 9, 8, 7 /* segment 1 */};

Template Parameters:

KeyInputIteratorItT – [inferred] Random-access input iterator over per-segment key-input iterators (may be a simple pointer type)
KeyOutputIteratorItT – [inferred] Random-access input iterator over per-segment key-output iterators (may be a simple pointer type)
ValueInputIteratorItT – [inferred] Random-access input iterator over per-segment value-input iterators (may be a simple pointer type)
ValueOutputIteratorItT – [inferred] Random-access input iterator over per-segment value-output iterators (may be a simple pointer type)
SegmentSizeParameterT – [inferred] Type of the segment_sizes argument
KParameterT – [inferred] Type of the k argument
NumSegmentsParameterT – [inferred] Type of the num_segments argument
EnvT – [inferred] Execution environment type. Default is cuda::std::execution::env<>.

Parameters:

d_keys_in – [in] Iterator such that d_keys_in[i] yields a random-access iterator to the keys of segment i
d_keys_out – [out] Iterator such that d_keys_out[i] yields a random-access output iterator for the top-k keys of segment i
d_values_in – [in] Iterator such that d_values_in[i] yields a random-access iterator to the values of segment i
d_values_out – [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
segment_sizes – [in] Annotated argument providing the per-segment sizes (e.g. cuda::args::constant<N> 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).
k – [in] The number of selected items per segment, given as a cuda::args annotation or a plain integral value.
num_segments – [in] The (uniform) number of segments, given as a cuda::args annotation or a plain integral value.
env – [in]
[optional] Execution environment. Must require determinism::not_guaranteed, tie_break::unspecified, and output_ordering::unsorted.

template<typename KeyInputIteratorItT, typename KeyOutputIteratorItT, typename ValueInputIteratorItT, typename ValueOutputIteratorItT, typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, typename EnvT = ::cuda::std::execution::env<>> static inline 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, const EnvT &env = {} )#

Finds, for each segment, the smallest K keys and their corresponding values from an unordered input sequence of key-value pairs.

Temporary storage for this operation. If d_temp_storage is nullptr, the required size is written to temp_storage_bytes without dereferencing iterators or launching kernels. Otherwise, d_temp_storage must point to a device-accessible allocation of at least temp_storage_bytes bytes. No special alignment is required. See Two-Phase API (explicit temporary storage management) for usage guidance.

A Simple Example#

constexpr int num_segments = 2;
constexpr int segment_size = 8;
constexpr int k            = 3;

auto keys_in    = thrust::device_vector<int>{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4};
auto keys_out   = thrust::device_vector<int>(num_segments * k, thrust::no_init);
auto values_out = thrust::device_vector<int>(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_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);

constexpr auto segment_sizes = cuda::args::constant<segment_size>{};
constexpr auto k_arg         = cuda::args::constant<k>{};
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::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<char> temp_storage(temp_storage_bytes, thrust::no_init);
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);
// 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<int>{-3, 1, 2, /* segment 0 */ 0, 1, 2 /* segment 1 */};

Note

The behavior is undefined if an output range overlaps another output range or any input range. Input ranges may overlap one another.

Template Parameters:

KeyInputIteratorItT – [inferred] Random-access input iterator over per-segment key-input iterators (may be a simple pointer type)
KeyOutputIteratorItT – [inferred] Random-access input iterator over per-segment key-output iterators (may be a simple pointer type)
ValueInputIteratorItT – [inferred] Random-access input iterator over per-segment value-input iterators (may be a simple pointer type)
ValueOutputIteratorItT – [inferred] Random-access input iterator over per-segment value-output iterators (may be a simple pointer type)
SegmentSizeParameterT – [inferred] Type of the segment_sizes argument
KParameterT – [inferred] Type of the k argument
NumSegmentsParameterT – [inferred] Type of the num_segments argument
EnvT – [inferred] Execution environment type. Default is cuda::std::execution::env<>.

Parameters:

d_temp_storage – [in] Device-accessible allocation of temporary storage. When nullptr, the required allocation size is written to temp_storage_bytes and no work is done.
temp_storage_bytes – [inout] Reference to size in bytes of d_temp_storage allocation
d_keys_in – [in] Iterator such that d_keys_in[i] yields a random-access iterator to the keys of segment i
d_keys_out – [out] Iterator such that d_keys_out[i] yields a random-access output iterator for the top-k keys of segment i
d_values_in – [in] Iterator such that d_values_in[i] yields a random-access iterator to the values of segment i
d_values_out – [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
segment_sizes – [in] Annotated argument providing the per-segment sizes (e.g. cuda::args::constant<N> 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).
k – [in] The number of selected items per segment, given as a cuda::args annotation or a plain integral value.
num_segments – [in] The (uniform) number of segments, given as a cuda::args annotation or a plain integral value.
env – [in]
[optional] Execution environment. Must require determinism::not_guaranteed, tie_break::unspecified, and output_ordering::unsorted.

template<typename KeyInputIteratorItT, typename KeyOutputIteratorItT, typename ValueInputIteratorItT, typename ValueOutputIteratorItT, typename SegmentSizeParameterT, typename KParameterT, typename NumSegmentsParameterT, typename EnvT = ::cuda::std::execution::env<>> static inline 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, const EnvT &env = {} )#

Finds, for each segment, the smallest K keys and their corresponding values. Environment-based overload that allocates temporary storage internally.

Snippet#

constexpr int num_segments = 2;
constexpr int segment_size = 8;
constexpr int k            = 3;

auto keys_in    = thrust::device_vector<int>{5, -3, 1, 7, 8, 2, 4, 6, /**/ 0, 9, 3, 2, 1, 8, 7, 4};
auto keys_out   = thrust::device_vector<int>(num_segments * k, thrust::no_init);
auto values_out = thrust::device_vector<int>(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_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);

cuda::stream stream{cuda::devices[0]};
auto env = cuda::std::execution::env{
  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(
  d_keys_in,
  d_keys_out,
  d_values_in,
  d_values_out,
  cuda::args::constant<segment_size>{},
  cuda::args::constant<k>{},
  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<int>{-3, 1, 2, /* segment 0 */ 0, 1, 2 /* segment 1 */};

Template Parameters:

KeyInputIteratorItT – [inferred] Random-access input iterator over per-segment key-input iterators (may be a simple pointer type)
KeyOutputIteratorItT – [inferred] Random-access input iterator over per-segment key-output iterators (may be a simple pointer type)
ValueInputIteratorItT – [inferred] Random-access input iterator over per-segment value-input iterators (may be a simple pointer type)
ValueOutputIteratorItT – [inferred] Random-access input iterator over per-segment value-output iterators (may be a simple pointer type)
SegmentSizeParameterT – [inferred] Type of the segment_sizes argument
KParameterT – [inferred] Type of the k argument
NumSegmentsParameterT – [inferred] Type of the num_segments argument
EnvT – [inferred] Execution environment type. Default is cuda::std::execution::env<>.

Parameters:

d_keys_in – [in] Iterator such that d_keys_in[i] yields a random-access iterator to the keys of segment i
d_keys_out – [out] Iterator such that d_keys_out[i] yields a random-access output iterator for the top-k keys of segment i
d_values_in – [in] Iterator such that d_values_in[i] yields a random-access iterator to the values of segment i
d_values_out – [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
segment_sizes – [in] Annotated argument providing the per-segment sizes (e.g. cuda::args::constant<N> 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).
k – [in] The number of selected items per segment, given as a cuda::args annotation or a plain integral value.
num_segments – [in] The (uniform) number of segments, given as a cuda::args annotation or a plain integral value.
env – [in]
[optional] Execution environment. Must require determinism::not_guaranteed, tie_break::unspecified, and output_ordering::unsorted.