dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity
Files
8d75652949948da1abc3e41fccbf0828638623a6
project_6/cccl_upstream/cub/cub/device/device_batched_topk.cuh
EngineX CI 56fd68e7dd [INFRA] Import NVIDIA/CCCL upstream as optimization reference library 
CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0
2026-07-30 09:35:51 +00:00

1133 lines
51 KiB
Plaintext
Raw Blame History

// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//! @file
//! cub::DeviceBatchedTopK provides device-wide, parallel operations for finding the K largest (or smallest) items
//! from many (small) segments of unordered data items residing within device-accessible memory.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/env_dispatch.cuh>
#include <cub/device/dispatch/dispatch_batched_topk.cuh>
#include <cub/device/dispatch/dispatch_common.cuh> // topk::select::{min, max}
#include <cub/util_type.cuh>
#include <cuda/__execution/determinism.h>
#include <cuda/__execution/output_ordering.h>
#include <cuda/__execution/require.h>
#include <cuda/__execution/tie_break.h>
#include <cuda/__execution/tune.h>
#include <cuda/__functional/call_or.h>
#include <cuda/__stream/get_stream.h>
#include <cuda/argument>
#include <cuda/std/__execution/env.h>
#include <cuda/std/__type_traits/is_integral.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__utility/move.h>
#include <cuda/std/cstdint>
#include <cuda/std/limits>
CUB_NAMESPACE_BEGIN
namespace detail
{
//! @cond
//! Shared implementation for all cub::DeviceBatchedTopK entry points.
//!
//! Validates the requested execution requirements and argument annotations, resolves the (optionally tuned) policy
//! selector from the environment, and forwards to the internal batched top-k dispatch. The selection direction is
//! threaded through as a compile-time `cuda::args::constant<SelectDirection>` so the kernel only emits the
//! requested (max OR min) code path.
//!
//! All current API-surface constraints are surfaced here as `static_assert`s so the diagnostic appears at the
//! `cub::DeviceBatchedTopK` call site rather than deep inside the kernel/agent instantiation.
template <topk::select SelectDirection,
typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename ValueInputIteratorItT,
typename ValueOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
ValueInputIteratorItT d_values_in,
ValueOutputIteratorItT d_values_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env)
{
// ---------------------------------------------------------------------------
// Execution requirements.
//
// Two orthogonal concerns govern the result: *which* items are selected (possibly refined by
// a tie-break preference) and the order in which they are written (output ordering). The committed default contract
// is the most reproducible behavior (determinism::gpu_to_gpu + tie_break::prefer_smaller_index +
// output_ordering::stable_sorted). Three rules are validated here:
// 1. determinism and tie_break must be acknowledged together (both specified, or both omitted default)
// 2. an explicit tie_break of prefer_smaller_index / prefer_larger_index fully pins the result set across GPUs and
// therefore requires determinism::gpu_to_gpu (it cannot be paired with run_to_run or not_guaranteed)
// 3. this initial API surface only implements the fully opted-out configuration (non-deterministic, unsorted).
// ---------------------------------------------------------------------------
static_assert(!::cuda::std::execution::__queryable_with<EnvT, ::cuda::execution::determinism::__get_determinism_t>,
"Determinism should be used inside cuda::execution::require to have an effect.");
static_assert(!::cuda::std::execution::__queryable_with<EnvT, ::cuda::execution::tie_break::__get_tie_break_t>,
"Tie-break should be used inside cuda::execution::require to have an effect.");
using requirements_t = ::cuda::std::execution::
__query_result_or_t<EnvT, ::cuda::execution::__get_requirements_t, ::cuda::std::execution::env<>>;
constexpr bool determinism_specified =
::cuda::std::execution::__queryable_with<requirements_t, ::cuda::execution::determinism::__get_determinism_t>;
constexpr bool tie_break_specified =
::cuda::std::execution::__queryable_with<requirements_t, ::cuda::execution::tie_break::__get_tie_break_t>;
using requested_determinism_t =
::cuda::std::execution::__query_result_or_t<requirements_t,
::cuda::execution::determinism::__get_determinism_t,
::cuda::execution::determinism::gpu_to_gpu_t>;
using requested_tie_break_t =
::cuda::std::execution::__query_result_or_t<requirements_t,
::cuda::execution::tie_break::__get_tie_break_t,
::cuda::execution::tie_break::prefer_smaller_index_t>;
using requested_order_t =
::cuda::std::execution::__query_result_or_t<requirements_t,
::cuda::execution::output_ordering::__get_output_ordering_t,
::cuda::execution::output_ordering::stable_sorted_t>;
constexpr bool determinism_and_tie_break_paired = (determinism_specified == tie_break_specified);
// Encodes rule 2 as the implication "a concrete tie-break requires gpu_to_gpu". The expression is the "or" form
// and satisfied in the two cases that are allowed: the tie-break is unspecified or the determinism is already
// gpu_to_gpu (which accepts any tie-break).
constexpr bool tie_break_compatible_with_determinism =
::cuda::std::is_same_v<requested_tie_break_t, ::cuda::execution::tie_break::unspecified_t>
|| ::cuda::std::is_same_v<requested_determinism_t, ::cuda::execution::determinism::gpu_to_gpu_t>;
constexpr bool is_non_deterministic_unsorted =
::cuda::std::is_same_v<requested_determinism_t, ::cuda::execution::determinism::not_guaranteed_t>
&& ::cuda::std::is_same_v<requested_order_t, ::cuda::execution::output_ordering::unsorted_t>;
static_assert(determinism_and_tie_break_paired,
"cub::DeviceBatchedTopK: determinism and tie_break requirements must be acknowledged together. Either "
"omit both to accept the defaults (cuda::execution::determinism::gpu_to_gpu and "
"cuda::execution::tie_break::prefer_smaller_index), or pass both explicitly inside "
"cuda::execution::require(...).");
static_assert(!determinism_and_tie_break_paired || tie_break_compatible_with_determinism,
"cub::DeviceBatchedTopK: a tie_break of cuda::execution::tie_break::prefer_smaller_index or "
"prefer_larger_index pins the result set across GPUs and therefore requires "
"cuda::execution::determinism::gpu_to_gpu (it cannot be combined with run_to_run or not_guaranteed).");
static_assert(
!determinism_and_tie_break_paired || !tie_break_compatible_with_determinism || is_non_deterministic_unsorted,
"cub::DeviceBatchedTopK currently only implements non-deterministic, unsorted output. Request it "
"explicitly with cuda::execution::require(cuda::execution::determinism::not_guaranteed, "
"cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted).");
// ---------------------------------------------------------------------------
// Resolve the (optionally tuned) policy selector from the environment.
// ---------------------------------------------------------------------------
using key_t = cub::detail::it_value_t<cub::detail::it_value_t<KeyInputIteratorItT>>;
using value_t = cub::detail::it_value_t<cub::detail::it_value_t<ValueInputIteratorItT>>;
using default_policy_selector_t = batched_topk::
policy_selector_from_types<key_t, value_t, ::cuda::std::int64_t, ::cuda::args::__traits<KParameterT>::highest>;
using tuning_env_t =
::cuda::__call_result_or_t<::cuda::execution::__get_tuning_t, ::cuda::std::execution::env<>, EnvT>;
using policy_selector_t = ::cuda::std::execution::
__query_result_or_t<tuning_env_t, batched_topk::batched_topk_policy, default_policy_selector_t>;
// ---------------------------------------------------------------------------
// Argument-annotation constraints surfaced at the call site.
// ---------------------------------------------------------------------------
static_assert(::cuda::args::__traits<NumSegmentsParameterT>::is_single_value,
"cub::DeviceBatchedTopK currently requires a single (uniform) number of segments resolved on the "
"host; pass num_segments as a single-value annotation (e.g. cuda::args::constant or "
"cuda::args::immediate), not a per-segment sequence.");
static_assert(
::cuda::args::__is_wrapper_v<SegmentSizeParameterT> || ::cuda::std::is_integral_v<SegmentSizeParameterT>,
"cub::DeviceBatchedTopK: segment_sizes must be a cuda::args annotation or a plain integral value "
"(taken as a uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap "
"per-segment sizes in cuda::args::deferred_sequence, or a single device-side value in "
"cuda::args::deferred.");
static_assert(::cuda::args::__is_wrapper_v<KParameterT> || ::cuda::std::is_integral_v<KParameterT>,
"cub::DeviceBatchedTopK: k must be a cuda::args annotation or a plain integral value (taken as a "
"uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap a per-segment k "
"in cuda::args::deferred_sequence, or a single device-side value in cuda::args::deferred.");
static_assert(
::cuda::args::__is_wrapper_v<NumSegmentsParameterT> || ::cuda::std::is_integral_v<NumSegmentsParameterT>,
"cub::DeviceBatchedTopK: num_segments must be a cuda::args annotation or a plain integral value. A "
"raw pointer or iterator is not accepted.");
const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);
// The total-number-of-items guarantee is intentionally not part of the initial public API surface. The dispatch
// only uses its element type to size internal large-segment offsets (the value itself is unused), so we pass a
// conservative 64-bit upper bound here.
constexpr auto total_num_items = ::cuda::args::immediate{::cuda::std::numeric_limits<::cuda::std::int64_t>::max()};
return batched_topk::dispatch(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
segment_sizes,
k,
::cuda::args::constant<SelectDirection>{},
num_segments,
total_num_items,
stream.get(),
policy_selector_t{});
}
//! @endcond
} // namespace detail
//! @rst
//! DeviceBatchedTopK provides device-wide, parallel operations for finding the largest (or smallest) K items from
//! many segments of unordered data items residing within device-accessible memory.
//!
//! .. versionadded:: 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.
//!
//! .. code-block:: c++
//!
//! // segment_sizes (k is analogous):
//! cuda::args::constant<256>{}; // fixed at compile time
//! cuda::args::immediate{n, cuda::args::bounds<1, 1024>()}; // host value, at most 1024
//! cuda::args::deferred_sequence{d_sizes, cuda::args::bounds<1, 1024>()}; // per-segment, each at most 1024
//!
//! // a 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 :ref:`cub-topk-requirements`) 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 :cpp:struct:`cub::DeviceTopK`, the result of ``DeviceBatchedTopK`` is governed by two orthogonal execution
//! requirements: *which* items are selected per segment (``cuda::execution::determinism``, optionally refined by
//! ``cuda::execution::tie_break``) and the order in which they are written (``cuda::execution::output_ordering``).
//! When the caller does not opt out, the committed default is the most reproducible behavior: deterministic results
//! (``cuda::execution::determinism::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 :ref:`cub-topk-requirements` for the full requirement model, worked examples, and guidance on choosing
//! requirements.
//!
//! .. 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.
//!
//! Usage Considerations
//! ++++++++++++++++++++++++++
//!
//! @cdp_class{DeviceBatchedTopK}
//!
//! @endrst
struct DeviceBatchedTopK
{
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds, for each segment, the largest K keys from an unordered input sequence of keys.
//!
//! .. note::
//!
//! The behavior is undefined if an output range overlaps another output range or any input range.
//! Input ranges may overlap one another.
//!
//! - @devicestorage
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin batched-topk-max-keys
//! :end-before: example-end batched-topk-max-keys
//!
//! @endrst
//!
//! @tparam KeyInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator
//!
//! @tparam KeyOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator
//!
//! @tparam SegmentSizeParameterT
//! **[inferred]** Type of the ``segment_sizes`` argument
//!
//! @tparam KParameterT
//! **[inferred]** Type of the ``k`` argument
//!
//! @tparam NumSegmentsParameterT
//! **[inferred]** Type of the ``num_segments`` argument
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_temp_storage
//! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to
//! `temp_storage_bytes` and no work is done.
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i`
//!
//! @param[out] d_keys_out
//! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i`
//!
//! @param[in] segment_sizes
//! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant<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).
//!
//! @param[in] k
//! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] num_segments
//! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Must require `determinism::not_guaranteed`,
//! `tie_break::unspecified`, and `output_ordering::unsorted`.
//! @endrst
template <typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT = ::cuda::std::execution::env<>>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MaxKeys");
return detail::dispatch_batched_topk<detail::topk::select::max>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
static_cast<NullType**>(nullptr),
static_cast<NullType**>(nullptr),
segment_sizes,
k,
num_segments,
::cuda::std::move(env));
}
//! @rst
//! Finds, for each segment, the largest K keys from an unordered input sequence of keys.
//!
//! This is an environment-based API that allocates and manages the required temporary storage internally using the
//! memory resource queried from the environment.
//!
//! .. note::
//!
//! The behavior is undefined if an output range overlaps another output range or any input range.
//! Input ranges may overlap one another.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin batched-topk-max-keys-env
//! :end-before: example-end batched-topk-max-keys-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator
//!
//! @tparam KeyOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator
//!
//! @tparam SegmentSizeParameterT
//! **[inferred]** Type of the ``segment_sizes`` argument
//!
//! @tparam KParameterT
//! **[inferred]** Type of the ``k`` argument
//!
//! @tparam NumSegmentsParameterT
//! **[inferred]** Type of the ``num_segments`` argument
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i`
//!
//! @param[out] d_keys_out
//! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i`
//!
//! @param[in] segment_sizes
//! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant<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).
//!
//! @param[in] k
//! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] num_segments
//! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Must require `determinism::not_guaranteed`,
//! `tie_break::unspecified`, and `output_ordering::unsorted`.
//! @endrst
template <typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT = ::cuda::std::execution::env<>>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys(
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxKeys");
return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) {
return detail::dispatch_batched_topk<detail::topk::select::max>(
storage,
bytes,
d_keys_in,
d_keys_out,
static_cast<NullType**>(nullptr),
static_cast<NullType**>(nullptr),
segment_sizes,
k,
num_segments,
env);
});
}
//! @rst
//! Finds, for each segment, the smallest K keys from an unordered input sequence of keys.
//!
//! .. note::
//!
//! The behavior is undefined if an output range overlaps another output range or any input range.
//! Input ranges may overlap one another.
//!
//! - @devicestorage
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin batched-topk-min-keys
//! :end-before: example-end batched-topk-min-keys
//!
//! @endrst
//!
//! @tparam KeyInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator
//!
//! @tparam KeyOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator
//!
//! @tparam SegmentSizeParameterT
//! **[inferred]** Type of the ``segment_sizes`` argument
//!
//! @tparam KParameterT
//! **[inferred]** Type of the ``k`` argument
//!
//! @tparam NumSegmentsParameterT
//! **[inferred]** Type of the ``num_segments`` argument
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_temp_storage
//! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to
//! `temp_storage_bytes` and no work is done.
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i`
//!
//! @param[out] d_keys_out
//! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i`
//!
//! @param[in] segment_sizes
//! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant<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).
//!
//! @param[in] k
//! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] num_segments
//! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Must require `determinism::not_guaranteed`,
//! `tie_break::unspecified`, and `output_ordering::unsorted`.
//! @endrst
template <typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT = ::cuda::std::execution::env<>>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinKeys(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MinKeys");
return detail::dispatch_batched_topk<detail::topk::select::min>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
static_cast<NullType**>(nullptr),
static_cast<NullType**>(nullptr),
segment_sizes,
k,
num_segments,
::cuda::std::move(env));
}
//! @rst
//! Finds, for each segment, the smallest K keys from an unordered input sequence of keys. Environment-based overload
//! that allocates temporary storage internally.
//!
//! .. note::
//!
//! The behavior is undefined if an output range overlaps another output range or any input range.
//! Input ranges may overlap one another.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin batched-topk-min-keys-env
//! :end-before: example-end batched-topk-min-keys-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator
//!
//! @tparam KeyOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator
//!
//! @tparam SegmentSizeParameterT
//! **[inferred]** Type of the ``segment_sizes`` argument
//!
//! @tparam KParameterT
//! **[inferred]** Type of the ``k`` argument
//!
//! @tparam NumSegmentsParameterT
//! **[inferred]** Type of the ``num_segments`` argument
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i`
//!
//! @param[out] d_keys_out
//! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i`
//!
//! @param[in] segment_sizes
//! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant<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).
//!
//! @param[in] k
//! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] num_segments
//! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Must require `determinism::not_guaranteed`,
//! `tie_break::unspecified`, and `output_ordering::unsorted`.
//! @endrst
template <typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT = ::cuda::std::execution::env<>>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinKeys(
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinKeys");
return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) {
return detail::dispatch_batched_topk<detail::topk::select::min>(
storage,
bytes,
d_keys_in,
d_keys_out,
static_cast<NullType**>(nullptr),
static_cast<NullType**>(nullptr),
segment_sizes,
k,
num_segments,
env);
});
}
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds, for each segment, the largest K keys and their corresponding values from an unordered input sequence of
//! key-value pairs.
//!
//! .. note::
//!
//! The behavior is undefined if an output range overlaps another output range or any input range.
//! Input ranges may overlap one another.
//!
//! - @devicestorage
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin batched-topk-max-pairs
//! :end-before: example-end batched-topk-max-pairs
//!
//! @endrst
//!
//! @tparam KeyInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator
//!
//! @tparam KeyOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator
//!
//! @tparam ValueInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment value-input iterators @iterator
//!
//! @tparam ValueOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment value-output iterators @iterator
//!
//! @tparam SegmentSizeParameterT
//! **[inferred]** Type of the ``segment_sizes`` argument
//!
//! @tparam KParameterT
//! **[inferred]** Type of the ``k`` argument
//!
//! @tparam NumSegmentsParameterT
//! **[inferred]** Type of the ``num_segments`` argument
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_temp_storage
//! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to
//! `temp_storage_bytes` and no work is done.
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i`
//!
//! @param[out] d_keys_out
//! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i`
//!
//! @param[in] d_values_in
//! Iterator such that `d_values_in[i]` yields a random-access iterator to the values of segment `i`
//!
//! @param[out] d_values_out
//! Iterator such that `d_values_out[i]` yields a random-access output iterator for the values corresponding to the
//! top-k keys of segment `i`
//!
//! @param[in] segment_sizes
//! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant<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).
//!
//! @param[in] k
//! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] num_segments
//! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Must require `determinism::not_guaranteed`,
//! `tie_break::unspecified`, and `output_ordering::unsorted`.
//! @endrst
template <typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename ValueInputIteratorItT,
typename ValueOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT = ::cuda::std::execution::env<>>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
ValueInputIteratorItT d_values_in,
ValueOutputIteratorItT d_values_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MaxPairs");
return detail::dispatch_batched_topk<detail::topk::select::max>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
segment_sizes,
k,
num_segments,
::cuda::std::move(env));
}
//! @rst
//! Finds, for each segment, the largest K keys and their corresponding values. Environment-based overload that
//! allocates temporary storage internally.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin batched-topk-max-pairs-env
//! :end-before: example-end batched-topk-max-pairs-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator
//!
//! @tparam KeyOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator
//!
//! @tparam ValueInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment value-input iterators @iterator
//!
//! @tparam ValueOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment value-output iterators @iterator
//!
//! @tparam SegmentSizeParameterT
//! **[inferred]** Type of the ``segment_sizes`` argument
//!
//! @tparam KParameterT
//! **[inferred]** Type of the ``k`` argument
//!
//! @tparam NumSegmentsParameterT
//! **[inferred]** Type of the ``num_segments`` argument
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i`
//!
//! @param[out] d_keys_out
//! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i`
//!
//! @param[in] d_values_in
//! Iterator such that `d_values_in[i]` yields a random-access iterator to the values of segment `i`
//!
//! @param[out] d_values_out
//! Iterator such that `d_values_out[i]` yields a random-access output iterator for the values corresponding to the
//! top-k keys of segment `i`
//!
//! @param[in] segment_sizes
//! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant<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).
//!
//! @param[in] k
//! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] num_segments
//! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Must require `determinism::not_guaranteed`,
//! `tie_break::unspecified`, and `output_ordering::unsorted`.
//! @endrst
template <typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename ValueInputIteratorItT,
typename ValueOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT = ::cuda::std::execution::env<>>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs(
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
ValueInputIteratorItT d_values_in,
ValueOutputIteratorItT d_values_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MaxPairs");
return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) {
return detail::dispatch_batched_topk<detail::topk::select::max>(
storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env);
});
}
//! @rst
//! Finds, for each segment, the smallest K keys and their corresponding values from an unordered input sequence of
//! key-value pairs.
//!
//! .. note::
//!
//! The behavior is undefined if an output range overlaps another output range or any input range.
//! Input ranges may overlap one another.
//!
//! - @devicestorage
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin batched-topk-min-pairs
//! :end-before: example-end batched-topk-min-pairs
//!
//! @endrst
//!
//! @tparam KeyInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator
//!
//! @tparam KeyOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator
//!
//! @tparam ValueInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment value-input iterators @iterator
//!
//! @tparam ValueOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment value-output iterators @iterator
//!
//! @tparam SegmentSizeParameterT
//! **[inferred]** Type of the ``segment_sizes`` argument
//!
//! @tparam KParameterT
//! **[inferred]** Type of the ``k`` argument
//!
//! @tparam NumSegmentsParameterT
//! **[inferred]** Type of the ``num_segments`` argument
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_temp_storage
//! Device-accessible allocation of temporary storage. When `nullptr`, the required allocation size is written to
//! `temp_storage_bytes` and no work is done.
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i`
//!
//! @param[out] d_keys_out
//! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i`
//!
//! @param[in] d_values_in
//! Iterator such that `d_values_in[i]` yields a random-access iterator to the values of segment `i`
//!
//! @param[out] d_values_out
//! Iterator such that `d_values_out[i]` yields a random-access output iterator for the values corresponding to the
//! top-k keys of segment `i`
//!
//! @param[in] segment_sizes
//! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant<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).
//!
//! @param[in] k
//! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] num_segments
//! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Must require `determinism::not_guaranteed`,
//! `tie_break::unspecified`, and `output_ordering::unsorted`.
//! @endrst
template <typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename ValueInputIteratorItT,
typename ValueOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT = ::cuda::std::execution::env<>>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinPairs(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
ValueInputIteratorItT d_values_in,
ValueOutputIteratorItT d_values_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceBatchedTopK::MinPairs");
return detail::dispatch_batched_topk<detail::topk::select::min>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
segment_sizes,
k,
num_segments,
::cuda::std::move(env));
}
//! @rst
//! Finds, for each segment, the smallest K keys and their corresponding values. Environment-based overload that
//! allocates temporary storage internally.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_batched_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin batched-topk-min-pairs-env
//! :end-before: example-end batched-topk-min-pairs-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-input iterators @iterator
//!
//! @tparam KeyOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment key-output iterators @iterator
//!
//! @tparam ValueInputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment value-input iterators @iterator
//!
//! @tparam ValueOutputIteratorItT
//! **[inferred]** Random-access input iterator over per-segment value-output iterators @iterator
//!
//! @tparam SegmentSizeParameterT
//! **[inferred]** Type of the ``segment_sizes`` argument
//!
//! @tparam KParameterT
//! **[inferred]** Type of the ``k`` argument
//!
//! @tparam NumSegmentsParameterT
//! **[inferred]** Type of the ``num_segments`` argument
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Iterator such that `d_keys_in[i]` yields a random-access iterator to the keys of segment `i`
//!
//! @param[out] d_keys_out
//! Iterator such that `d_keys_out[i]` yields a random-access output iterator for the top-k keys of segment `i`
//!
//! @param[in] d_values_in
//! Iterator such that `d_values_in[i]` yields a random-access iterator to the values of segment `i`
//!
//! @param[out] d_values_out
//! Iterator such that `d_values_out[i]` yields a random-access output iterator for the values corresponding to the
//! top-k keys of segment `i`
//!
//! @param[in] segment_sizes
//! Annotated argument providing the per-segment sizes (e.g. `cuda::args::constant<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).
//!
//! @param[in] k
//! The number of selected items per segment, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] num_segments
//! The (uniform) number of segments, given as a `cuda::args` annotation or a plain integral value.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Must require `determinism::not_guaranteed`,
//! `tie_break::unspecified`, and `output_ordering::unsorted`.
//! @endrst
template <typename KeyInputIteratorItT,
typename KeyOutputIteratorItT,
typename ValueInputIteratorItT,
typename ValueOutputIteratorItT,
typename SegmentSizeParameterT,
typename KParameterT,
typename NumSegmentsParameterT,
typename EnvT = ::cuda::std::execution::env<>>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinPairs(
KeyInputIteratorItT d_keys_in,
KeyOutputIteratorItT d_keys_out,
ValueInputIteratorItT d_values_in,
ValueOutputIteratorItT d_values_out,
SegmentSizeParameterT segment_sizes,
KParameterT k,
NumSegmentsParameterT num_segments,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceBatchedTopK::MinPairs");
return detail::dispatch_with_env(env, [&](auto /* tuning */, void* storage, size_t& bytes, auto /* stream */) {
return detail::dispatch_batched_topk<detail::topk::select::min>(
storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env);
});
}
};
CUB_NAMESPACE_END
Reference in New Issue View Git Blame Copy Permalink