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. A plain integral (no bound) is still subject to each parameter’s constraints: for segment_sizes in particular, a plain signed integral such as int is rejected because its maximum exceeds the supported maximum segment size of 2^21 (see Which form each parameter accepts and Current constraints below).

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 known on the host (constant, immediate, or a plain integral); a negative count is treated as no work (an empty batch). segment_sizes must have a statically-known maximum not exceeding the supported 2^21 (about 2 million; see Current constraints below): a type whose maximum already fits (a narrow type such as uint8_t, int16_t, or uint16_t) is accepted without an explicit bound, while a type whose maximum exceeds 2^21 (e.g. int32_t, uint32_t, or int64_t) must carry a compile-time upper bound (a constant<N> or cuda::args::bounds<lo, hi>()). A negative statically-known lower bound is allowed: negative runtime sizes are clamped to an empty segment (size 0). A non-negative lower bound is trusted – passing an actual value outside its declared bound (for instance a negative value under a non-negative bound) is a caller precondition violation (undefined behavior). k has no algorithm-imposed maximum: a k larger than a segment’s size selects that whole segment. Its lower bound follows the same rule as segment_sizes above – under a negative statically-known lower bound a negative runtime k is clamped to 0 (selecting nothing), while under a non-negative lower bound a negative value is a caller precondition violation (undefined behavior). 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 compile error is emitted if violated):

  • Host-only. Unlike most CUB algorithms, DeviceBatchedTopK does not support CUDA dynamic parallelism: its methods must be invoked from host code, not from device code.

  • Segment size is architecture-dependent. On pre-Hopper GPUs (compute capability < 9.0) 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. On Hopper and newer GPUs (compute capability >= 9.0) the thread-block-cluster backend also handles larger segments that exceed this per-block limit. Both uniform (fixed) and variable segment sizes are supported. Independent of the architecture – and independent of the integer type used for the segment_sizes argument (a wider type such as int64_t does not raise it) – an individual segment is currently limited to a maximum of 2^21 (about 2 million) items, enforced at compile time from the statically-known maximum segment size. Larger segments are future work. A type whose maximum already lies within it (a narrow type such as uint8_t, int16_t, or uint16_t) is accepted un-annotated; a type whose maximum exceeds 2^21 (e.g. int32_t, uint32_t, or int64_t) must carry a compile-time cuda::args::bounds whose upper end does not exceed 2^21 (see Which form each parameter accepts above for how the lower bound and out-of-bound values are handled).

  • k is at most 64 bits wide. The element type of k may be no wider than 64 bits. The device clamps k to the segment size through a 64-bit intermediate, so a wider integer type (e.g. __int128) is rejected to avoid a silent wrap. k itself has no algorithm-imposed maximum (see Which form each parameter accepts above).

  • Uniform number of segments. num_segments must be a single value (constant, immediate, or a plain integral) resolved on the host, and must not exceed 2^31 - 1 (it sizes the launch grid). A negative count is treated as no work (an empty batch). A deferred (device-resident) count is not supported at this time.

  • Unsorted output required. Only cuda::execution::output_ordering::unsorted is implemented; the sorted orderings of the default contract described in Determinism, tie-breaking, and output ordering below (and hence an empty, no-requirement environment, which defaults to stable_sorted) are rejected at compile time. The supported determinism / tie_break requirements depend on the architecture – see the Current support note below. determinism and tie_break must always be specified together, or both omitted to take the default.

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.

Note

Current support. Only the unsorted output ordering is implemented; cuda::execution::output_ordering::unsorted must be requested explicitly (sorted / stable_sorted, and thus an empty, no-requirement environment, are rejected at compile time). The supported selection requirements and segment sizes differ by architecture:

  • Pre-Hopper (compute capability < 9.0): only the fully non-deterministic request (determinism::not_guaranteed, tie_break::unspecified) is supported, and every segment must fit a single thread block. Deterministic / tie-break requests and larger segments require SM 9.0+ and are diagnosed at compile time (or, when CUB_DISABLE_TOPK_UNSUPPORTED_ARCH_ASSERT is defined, deferred to runtime as cudaErrorNotSupported).

  • Hopper and newer (compute capability >= 9.0): all five acknowledged (determinism, tie_break) pairs are supported – (not_guaranteed, unspecified), (run_to_run, unspecified), and (gpu_to_gpu, {unspecified, prefer_smaller_index, prefer_larger_index}) – and segments larger than a single thread block are also supported.

When determinism::not_guaranteed is requested the per-segment output 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). Its statically-known maximum must lie within the currently supported range (see the Current constraints section for the value and its architecture dependence): narrow types (int8_t/int16_t/uint16_t) qualify unannotated, while wider types (int32_t/uint32_t/ int64_t) must carry a compile-time cuda::args::bounds. A negative lower bound is allowed; negative runtime sizes clamp to an empty segment. 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. It has no algorithm-imposed maximum; a k larger than a segment’s size selects that whole segment. Like segment_sizes, a negative lower bound is allowed and a negative runtime k is then clamped to 0 (selecting nothing).

  • num_segments – [in] The number of segments, given as a cuda::args annotation or a plain integral value.

  • env – [in]

    [optional] Execution environment. Must require output_ordering::unsorted (sorted / stable_sorted, and thus an empty environment, are not yet supported). The selection requirements may be any acknowledged (determinism, tie_break) pair: (not_guaranteed, unspecified), (run_to_run, unspecified), or gpu_to_gpu with unspecified / prefer_smaller_index / prefer_larger_index. Deterministic requests require SM 9.0+.

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). Its statically-known maximum must lie within the currently supported range (see the Current constraints section for the value and its architecture dependence): narrow types (int8_t/int16_t/uint16_t) qualify unannotated, while wider types (int32_t/uint32_t/ int64_t) must carry a compile-time cuda::args::bounds. A negative lower bound is allowed; negative runtime sizes clamp to an empty segment. 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. It has no algorithm-imposed maximum; a k larger than a segment’s size selects that whole segment. Like segment_sizes, a negative lower bound is allowed and a negative runtime k is then clamped to 0 (selecting nothing).

  • num_segments – [in] The number of segments, given as a cuda::args annotation or a plain integral value.

  • env – [in]

    [optional] Execution environment. Must require output_ordering::unsorted (sorted / stable_sorted, and thus an empty environment, are not yet supported). The selection requirements may be any acknowledged (determinism, tie_break) pair: (not_guaranteed, unspecified), (run_to_run, unspecified), or gpu_to_gpu with unspecified / prefer_smaller_index / prefer_larger_index. Deterministic requests require SM 9.0+.

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). Its statically-known maximum must lie within the currently supported range (see the Current constraints section for the value and its architecture dependence): narrow types (int8_t/int16_t/uint16_t) qualify unannotated, while wider types (int32_t/uint32_t/ int64_t) must carry a compile-time cuda::args::bounds. A negative lower bound is allowed; negative runtime sizes clamp to an empty segment. 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. It has no algorithm-imposed maximum; a k larger than a segment’s size selects that whole segment. Like segment_sizes, a negative lower bound is allowed and a negative runtime k is then clamped to 0 (selecting nothing).

  • num_segments – [in] The number of segments, given as a cuda::args annotation or a plain integral value.

  • env – [in]

    [optional] Execution environment. Must require output_ordering::unsorted (sorted / stable_sorted, and thus an empty environment, are not yet supported). The selection requirements may be any acknowledged (determinism, tie_break) pair: (not_guaranteed, unspecified), (run_to_run, unspecified), or gpu_to_gpu with unspecified / prefer_smaller_index / prefer_larger_index. Deterministic requests require SM 9.0+.

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). Its statically-known maximum must lie within the currently supported range (see the Current constraints section for the value and its architecture dependence): narrow types (int8_t/int16_t/uint16_t) qualify unannotated, while wider types (int32_t/uint32_t/ int64_t) must carry a compile-time cuda::args::bounds. A negative lower bound is allowed; negative runtime sizes clamp to an empty segment. 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. It has no algorithm-imposed maximum; a k larger than a segment’s size selects that whole segment. Like segment_sizes, a negative lower bound is allowed and a negative runtime k is then clamped to 0 (selecting nothing).

  • num_segments – [in] The number of segments, given as a cuda::args annotation or a plain integral value.

  • env – [in]

    [optional] Execution environment. Must require output_ordering::unsorted (sorted / stable_sorted, and thus an empty environment, are not yet supported). The selection requirements may be any acknowledged (determinism, tie_break) pair: (not_guaranteed, unspecified), (run_to_run, unspecified), or gpu_to_gpu with unspecified / prefer_smaller_index / prefer_larger_index. Deterministic requests require SM 9.0+.

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). Its statically-known maximum must lie within the currently supported range (see the Current constraints section for the value and its architecture dependence): narrow types (int8_t/int16_t/uint16_t) qualify unannotated, while wider types (int32_t/uint32_t/ int64_t) must carry a compile-time cuda::args::bounds. A negative lower bound is allowed; negative runtime sizes clamp to an empty segment. 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. It has no algorithm-imposed maximum; a k larger than a segment’s size selects that whole segment. Like segment_sizes, a negative lower bound is allowed and a negative runtime k is then clamped to 0 (selecting nothing).

  • num_segments – [in] The number of segments, given as a cuda::args annotation or a plain integral value.

  • env – [in]

    [optional] Execution environment. Must require output_ordering::unsorted (sorted / stable_sorted, and thus an empty environment, are not yet supported). The selection requirements may be any acknowledged (determinism, tie_break) pair: (not_guaranteed, unspecified), (run_to_run, unspecified), or gpu_to_gpu with unspecified / prefer_smaller_index / prefer_larger_index. Deterministic requests require SM 9.0+.

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). Its statically-known maximum must lie within the currently supported range (see the Current constraints section for the value and its architecture dependence): narrow types (int8_t/int16_t/uint16_t) qualify unannotated, while wider types (int32_t/uint32_t/ int64_t) must carry a compile-time cuda::args::bounds. A negative lower bound is allowed; negative runtime sizes clamp to an empty segment. 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. It has no algorithm-imposed maximum; a k larger than a segment’s size selects that whole segment. Like segment_sizes, a negative lower bound is allowed and a negative runtime k is then clamped to 0 (selecting nothing).

  • num_segments – [in] The number of segments, given as a cuda::args annotation or a plain integral value.

  • env – [in]

    [optional] Execution environment. Must require output_ordering::unsorted (sorted / stable_sorted, and thus an empty environment, are not yet supported). The selection requirements may be any acknowledged (determinism, tie_break) pair: (not_guaranteed, unspecified), (run_to_run, unspecified), or gpu_to_gpu with unspecified / prefer_smaller_index / prefer_larger_index. Deterministic requests require SM 9.0+.

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). Its statically-known maximum must lie within the currently supported range (see the Current constraints section for the value and its architecture dependence): narrow types (int8_t/int16_t/uint16_t) qualify unannotated, while wider types (int32_t/uint32_t/ int64_t) must carry a compile-time cuda::args::bounds. A negative lower bound is allowed; negative runtime sizes clamp to an empty segment. 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. It has no algorithm-imposed maximum; a k larger than a segment’s size selects that whole segment. Like segment_sizes, a negative lower bound is allowed and a negative runtime k is then clamped to 0 (selecting nothing).

  • num_segments – [in] The number of segments, given as a cuda::args annotation or a plain integral value.

  • env – [in]

    [optional] Execution environment. Must require output_ordering::unsorted (sorted / stable_sorted, and thus an empty environment, are not yet supported). The selection requirements may be any acknowledged (determinism, tie_break) pair: (not_guaranteed, unspecified), (run_to_run, unspecified), or gpu_to_gpu with unspecified / prefer_smaller_index / prefer_larger_index. Deterministic requests require SM 9.0+.

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). Its statically-known maximum must lie within the currently supported range (see the Current constraints section for the value and its architecture dependence): narrow types (int8_t/int16_t/uint16_t) qualify unannotated, while wider types (int32_t/uint32_t/ int64_t) must carry a compile-time cuda::args::bounds. A negative lower bound is allowed; negative runtime sizes clamp to an empty segment. 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. It has no algorithm-imposed maximum; a k larger than a segment’s size selects that whole segment. Like segment_sizes, a negative lower bound is allowed and a negative runtime k is then clamped to 0 (selecting nothing).

  • num_segments – [in] The number of segments, given as a cuda::args annotation or a plain integral value.

  • env – [in]

    [optional] Execution environment. Must require output_ordering::unsorted (sorted / stable_sorted, and thus an empty environment, are not yet supported). The selection requirements may be any acknowledged (determinism, tie_break) pair: (not_guaranteed, unspecified), (run_to_run, unspecified), or gpu_to_gpu with unspecified / prefer_smaller_index / prefer_larger_index. Deterministic requests require SM 9.0+.