dylanyunlong/project_6_89d52222
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity
Files
c1065aaf2c4fb5c593c801115ff79bced7a4a855
project_6_89d52222/cccl_upstream/cub/cub/device/device_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

2008 lines
72 KiB
Plaintext
Raw Blame History

// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//! @file
//! cub::DeviceTopK provides device-wide, parallel operations for finding the K largest (or smallest) items from
//! sequences of data
#pragma once
#include <cub/config.cuh>
#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK
# if _CCCL_COMPILER(NVRTC)
# error \
"Including <cub/device/device_topk.cuh> is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. <cub/block/block_reduce.cuh>). You can define CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK to disable this warning."
# endif // _CCCL_COMPILER(NVRTC)
#endif // CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK
#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/choose_offset.cuh>
#include <cub/detail/env_dispatch.cuh>
#include <cub/device/dispatch/dispatch_topk.cuh>
#include <cuda/__execution/determinism.h>
#include <cuda/__execution/output_ordering.h>
#include <cuda/__execution/require.h>
#include <cuda/__execution/tune.h>
#include <cuda/__functional/call_or.h>
#include <cuda/__stream/get_stream.h>
#include <cuda/std/__execution/env.h>
#include <cuda/std/__type_traits/enable_if.h>
#include <cuda/std/__utility/move.h>
CUB_NAMESPACE_BEGIN
namespace detail
{
template <topk::select SelectDirection,
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_topk(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env)
{
// Offset type selection
using offset_t = choose_offset_t<NumItemsT>;
using out_offset_t = ::cuda::std::
conditional_t<sizeof(offset_t) < sizeof(choose_offset_t<NumOutItemsT>), offset_t, choose_offset_t<NumOutItemsT>>;
// Query environment properties to determine if the user-requested configuration is supported
static_assert(!::cuda::std::execution::__queryable_with<EnvT, ::cuda::execution::determinism::__get_determinism_t>,
"Determinism should be used inside requires to have an effect.");
using requirements_t = ::cuda::std::execution::
__query_result_or_t<EnvT, ::cuda::execution::__get_requirements_t, ::cuda::std::execution::env<>>;
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_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 auto is_determinism_not_guaranteed =
::cuda::std::is_same_v<requested_determinism_t, ::cuda::execution::determinism::not_guaranteed_t>;
constexpr auto is_output_order_unsorted =
::cuda::std::is_same_v<requested_order_t, ::cuda::execution::output_ordering::unsorted_t>;
// We only support the case where determinism is not guaranteed and output order is unsorted
static_assert(is_determinism_not_guaranteed && is_output_order_unsorted,
"cub::DeviceTopK only supports the case where determinism is not guaranteed and output order is "
"unsorted.");
// TODO (elstehle): align requirement validation with cub::DeviceBatchedTopK in CCCL 4.0. cub::DeviceTopK does not
// yet inspect cuda::execution::tie_break, so it still accepts requirement combinations that cub::DeviceBatchedTopK
// rejects. It should enforce that determinism and tie_break are requested together (or both omitted to take the
// default) and that an explicit tie_break requires cuda::execution::determinism::gpu_to_gpu.
// Query relevant properties from the environment
auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env);
// Extract policy selector from environment tuning
using default_policy_selector_t = topk::policy_selector_from_types<it_value_t<KeyInputIteratorT>>;
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, topk::topk_policy, default_policy_selector_t>;
return topk::dispatch<SelectDirection>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
static_cast<offset_t>(num_items),
static_cast<out_offset_t>(k),
decomposer,
stream.get(),
policy_selector_t{});
}
template <topk::select SelectDirection,
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT>
CUB_RUNTIME_FUNCTION static cudaError_t dispatch_topk_hub(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env)
{
return dispatch_topk<SelectDirection>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
num_items,
k,
detail::identity_decomposer_t{},
env);
}
} // namespace detail
//! @rst
//! DeviceTopK provides device-wide, parallel operations for finding the largest (or smallest) K items from sequences of
//! unordered data items residing within device-accessible memory.
//!
//! Overview
//! ++++++++++++++++++++++++++
//!
//! The TopK algorithm tries to find the largest (or smallest) K items in an unordered list. A related problem is called
//! `K selection problem <https://en.wikipedia.org/wiki/Selection_algorithm>`_, which finds the Kth largest
//! (or smallest) values in a list.
//! DeviceTopK will return K items in an unspecified order as results. It is based on an algorithm called
//! `AIR TopK <https://dl.acm.org/doi/10.1145/3581784.3607062>`_.
//!
//! Supported Types
//! ++++++++++++++++++++++++++
//!
//! DeviceTopK can process all of the built-in C++ numeric primitive types (`unsigned char`, `int`, `double`, etc.) as
//! well as CUDA's `__half` and `__nv_bfloat16` 16-bit floating-point types. User-defined types are supported as long
//! as a decomposer object is provided.
//!
//! Determinism, tie-breaking, and output ordering
//! +++++++++++++++++++++++++++++++++++++++++++++++
//!
//! The result of ``DeviceTopK`` is governed by two orthogonal execution requirements: *which* items are selected
//! (``cuda::execution::determinism``, optionally refined by ``cuda::execution::tie_break``) and the order in which
//! they are written (``cuda::execution::output_ordering``). When the caller does not opt out, the committed default
//! is the most reproducible behavior: deterministic results (``cuda::execution::determinism::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.
//!
//! 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::output_ordering::unsorted)``. Any other combination (including an empty, no-requirement
//! environment) is rejected at compile time. In this configuration the 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{DeviceTopK}
//!
//! Performance
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! @linear_performance{top-k}
//!
//! @endrst
struct DeviceTopK
{
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds the largest K keys and their corresponding values from an unordered input sequence of key-value pairs.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! - @devicestorage
//!
//! .. versionadded:: 3.3.0
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! The following code snippet demonstrates how to use the `cub::DeviceTopK::MaxPairs` function to find the largest K
//! items:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-max-pairs-non-deterministic-unsorted
//! :end-before: example-end topk-max-pairs-non-deterministic-unsorted
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam ValueInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input values @iterator
//!
//! @tparam ValueOutputIteratorT
//! **[inferred]** Random-access input iterator type for writing output values @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @param[in] d_temp_storage
//! @devicestorage
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] d_values_in
//! Random-access iterator to the input sequence containing the values associated to each key
//!
//! @param[out] d_values_out
//! Random-access iterator to the output sequence of values, corresponding to the top k keys, where k values will be
//! written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in` and `d_values_in` each
//!
//! @param[in] k
//! The value of K, which is the number of largest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, EnvT>, int> = 0>
CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceTopK::MaxPairs");
return detail::dispatch_topk<detail::topk::select::max>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
num_items,
k,
detail::identity_decomposer_t{},
::cuda::std::move(env));
}
//! @rst
//! Finds the largest K keys and their corresponding values from an unordered input sequence of key-value pairs.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! .. versionadded:: 3.5.0
//! First appears in CUDA Toolkit 13.5.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! Unlike the temp-storage overload, this overload allocates and manages the required temporary
//! storage internally using the memory resource queried from the environment.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-max-pairs-env
//! :end-before: example-end topk-max-pairs-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam ValueInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input values @iterator
//!
//! @tparam ValueOutputIteratorT
//! **[inferred]** Random-access input iterator type for writing output values @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] d_values_in
//! Random-access iterator to the input sequence containing the values associated to each key
//!
//! @param[out] d_values_out
//! Random-access iterator to the output sequence of values, corresponding to the top k keys, where k values will be
//! written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in` and `d_values_in` each
//!
//! @param[in] k
//! The value of K, which is the number of largest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, EnvT>, int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs(
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceTopK::MaxPairs");
return detail::dispatch_with_env(
env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) {
return detail::dispatch_topk<detail::topk::select::max>(
storage,
bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
num_items,
k,
detail::identity_decomposer_t{},
env);
});
}
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds the largest K keys and their corresponding values from an unordered input sequence of key-value pairs,
//! using a decomposer to interpret user-defined key types.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! - @devicestorage
//!
//! .. versionadded:: 3.4.0
//! First appears in CUDA Toolkit 13.4.
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Let's consider a user-defined ``custom_t`` type below. To find the top-k
//! elements of an array of ``custom_t`` objects, we have to tell CUB about
//! relevant members of the ``custom_t`` type. We do this by providing a
//! decomposer that returns a tuple of references to relevant members of the key.
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-custom-type
//! :end-before: example-end topk-custom-type
//!
//! The following snippet shows how to find the top-k largest pairs of ``custom_t``
//! objects using ``cub::DeviceTopK::MaxPairs``:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-max-pairs-custom-type
//! :end-before: example-end topk-max-pairs-custom-type
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam ValueInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input values @iterator
//!
//! @tparam ValueOutputIteratorT
//! **[inferred]** Random-access input iterator type for writing output values @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam DecomposerT
//! **[inferred]** Type of a callable object responsible for decomposing a key into a tuple of references to its
//! constituent arithmetic types.
//!
//! @param[in] d_temp_storage
//! @devicestorage
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] d_values_in
//! Random-access iterator to the input sequence containing the values associated to each key
//!
//! @param[out] d_values_out
//! Random-access iterator to the output sequence of values, corresponding to the top k keys, where k values will be
//! written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in` and `d_values_in` each
//!
//! @param[in] k
//! The value of K, which is the number of largest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param decomposer
//! Callable object responsible for decomposing a key into a tuple of references to its constituent arithmetic
//! types.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`.
//! @endrst
template <typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT = ::cuda::std::execution::env<>>
CUB_RUNTIME_FUNCTION static //
::cuda::std::enable_if_t<detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, DecomposerT>,
cudaError_t>
MaxPairs(void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceTopK::MaxPairs");
using key_t = detail::it_value_t<KeyInputIteratorT>;
static_assert(!detail::radix::can_twiddle<key_t>,
"Custom decomposers are not supported for fundamental types; "
"use the non-decomposer API overload instead");
return detail::dispatch_topk<detail::topk::select::max>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
num_items,
k,
decomposer,
::cuda::std::move(env));
}
//! @rst
//! Finds the largest K keys and their corresponding values from an unordered input sequence of key-value pairs,
//! using a decomposer to interpret user-defined key types.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! .. versionadded:: 3.5.0
//! First appears in CUDA Toolkit 13.5.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! Unlike the temp-storage overload, this overload allocates and manages the required temporary
//! storage internally using the memory resource queried from the environment.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-max-pairs-decomposer-env
//! :end-before: example-end topk-max-pairs-decomposer-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam ValueInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input values @iterator
//!
//! @tparam ValueOutputIteratorT
//! **[inferred]** Random-access input iterator type for writing output values @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam DecomposerT
//! **[inferred]** Type of a callable object responsible for decomposing a key into a tuple of references to its
//! constituent arithmetic types.
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] d_values_in
//! Random-access iterator to the input sequence containing the values associated to each key
//!
//! @param[out] d_values_out
//! Random-access iterator to the output sequence of values, corresponding to the top k keys, where k values will be
//! written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in` and `d_values_in` each
//!
//! @param[in] k
//! The value of K, which is the number of largest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] decomposer
//! Callable object responsible for decomposing a key into a tuple of references to its constituent arithmetic
//! types.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, DecomposerT>,
int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxPairs(
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceTopK::MaxPairs");
using key_t = detail::it_value_t<KeyInputIteratorT>;
static_assert(!detail::radix::can_twiddle<key_t>,
"Custom decomposers are not supported for fundamental types; "
"use the non-decomposer API overload instead");
return detail::dispatch_with_env(
env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) {
return detail::dispatch_topk<detail::topk::select::max>(
storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, num_items, k, decomposer, env);
});
}
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds the lowest K keys and their corresponding values from an unordered input sequence of key-value pairs.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! - @devicestorage
//!
//! .. versionadded:: 3.3.0
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! The following code snippet demonstrates how to use the `cub::DeviceTopK::MinPairs` function to find the lowest K
//! items:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-min-pairs-non-deterministic-unsorted
//! :end-before: example-end topk-min-pairs-non-deterministic-unsorted
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam ValueInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input values @iterator
//!
//! @tparam ValueOutputIteratorT
//! **[inferred]** Random-access input iterator type for writing output values @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @param[in] d_temp_storage
//! @devicestorage
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] d_values_in
//! Random-access iterator to the input sequence containing the values associated to each key
//!
//! @param[out] d_values_out
//! Random-access iterator to the output sequence of values, corresponding to the top k keys, where k values will be
//! written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in` and `d_values_in` each
//!
//! @param[in] k
//! The value of K, which is the number of lowest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, EnvT>, int> = 0>
CUB_RUNTIME_FUNCTION static cudaError_t MinPairs(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceTopK::MinPairs");
return detail::dispatch_topk<detail::topk::select::min>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
num_items,
k,
detail::identity_decomposer_t{},
::cuda::std::move(env));
}
//! @rst
//! Finds the smallest K keys and their corresponding values from an unordered input sequence of key-value pairs.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! .. versionadded:: 3.5.0
//! First appears in CUDA Toolkit 13.5.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! Unlike the temp-storage overload, this overload allocates and manages the required temporary
//! storage internally using the memory resource queried from the environment.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-min-pairs-env
//! :end-before: example-end topk-min-pairs-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam ValueInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input values @iterator
//!
//! @tparam ValueOutputIteratorT
//! **[inferred]** Random-access input iterator type for writing output values @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] d_values_in
//! Random-access iterator to the input sequence containing the values associated to each key
//!
//! @param[out] d_values_out
//! Random-access iterator to the output sequence of values, corresponding to the top k keys, where k values will be
//! written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in` and `d_values_in` each
//!
//! @param[in] k
//! The value of K, which is the number of lowest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, EnvT>, int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinPairs(
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceTopK::MinPairs");
return detail::dispatch_with_env(
env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) {
return detail::dispatch_topk<detail::topk::select::min>(
storage,
bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
num_items,
k,
detail::identity_decomposer_t{},
env);
});
}
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds the lowest K keys and their corresponding values from an unordered input sequence of key-value pairs,
//! using a decomposer to interpret user-defined key types.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! - @devicestorage
//!
//! .. versionadded:: 3.4.0
//! First appears in CUDA Toolkit 13.4.
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Let's consider a user-defined ``custom_t`` type below. To find the top-k
//! elements of an array of ``custom_t`` objects, we have to tell CUB about
//! relevant members of the ``custom_t`` type. We do this by providing a
//! decomposer that returns a tuple of references to relevant members of the key.
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-custom-type
//! :end-before: example-end topk-custom-type
//!
//! The following snippet shows how to find the top-k smallest pairs of ``custom_t``
//! objects using ``cub::DeviceTopK::MinPairs``:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-min-pairs-custom-type
//! :end-before: example-end topk-min-pairs-custom-type
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam ValueInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input values @iterator
//!
//! @tparam ValueOutputIteratorT
//! **[inferred]** Random-access input iterator type for writing output values @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam DecomposerT
//! **[inferred]** Type of a callable object responsible for decomposing a key into a tuple of references to its
//! constituent arithmetic types.
//!
//! @param[in] d_temp_storage
//! @devicestorage
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] d_values_in
//! Random-access iterator to the input sequence containing the values associated to each key
//!
//! @param[out] d_values_out
//! Random-access iterator to the output sequence of values, corresponding to the top k keys, where k values will be
//! written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in` and `d_values_in` each
//!
//! @param[in] k
//! The value of K, which is the number of lowest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param decomposer
//! Callable object responsible for decomposing a key into a tuple of references to its constituent arithmetic
//! types.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`.
//! @endrst
template <typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT = ::cuda::std::execution::env<>>
CUB_RUNTIME_FUNCTION static //
::cuda::std::enable_if_t<detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, DecomposerT>,
cudaError_t>
MinPairs(void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceTopK::MinPairs");
using key_t = detail::it_value_t<KeyInputIteratorT>;
static_assert(!detail::radix::can_twiddle<key_t>,
"Custom decomposers are not supported for fundamental types; "
"use the non-decomposer API overload instead");
return detail::dispatch_topk<detail::topk::select::min>(
d_temp_storage,
temp_storage_bytes,
d_keys_in,
d_keys_out,
d_values_in,
d_values_out,
num_items,
k,
decomposer,
::cuda::std::move(env));
}
//! @rst
//! Finds the smallest K keys and their corresponding values from an unordered input sequence of key-value pairs,
//! using a decomposer to interpret user-defined key types.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! .. versionadded:: 3.5.0
//! First appears in CUDA Toolkit 13.5.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! Unlike the temp-storage overload, this overload allocates and manages the required temporary
//! storage internally using the memory resource queried from the environment.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-min-pairs-decomposer-env
//! :end-before: example-end topk-min-pairs-decomposer-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam ValueInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input values @iterator
//!
//! @tparam ValueOutputIteratorT
//! **[inferred]** Random-access input iterator type for writing output values @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam DecomposerT
//! **[inferred]** Type of a callable object responsible for decomposing a key into a tuple of references to its
//! constituent arithmetic types.
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] d_values_in
//! Random-access iterator to the input sequence containing the values associated to each key
//!
//! @param[out] d_values_out
//! Random-access iterator to the output sequence of values, corresponding to the top k keys, where k values will be
//! written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in` and `d_values_in` each
//!
//! @param[in] k
//! The value of K, which is the number of lowest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] decomposer
//! Callable object responsible for decomposing a key into a tuple of references to its constituent arithmetic
//! types.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename ValueInputIteratorT,
typename ValueOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, DecomposerT>,
int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinPairs(
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
ValueInputIteratorT d_values_in,
ValueOutputIteratorT d_values_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceTopK::MinPairs");
using key_t = detail::it_value_t<KeyInputIteratorT>;
static_assert(!detail::radix::can_twiddle<key_t>,
"Custom decomposers are not supported for fundamental types; "
"use the non-decomposer API overload instead");
return detail::dispatch_with_env(
env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) {
return detail::dispatch_topk<detail::topk::select::min>(
storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, num_items, k, decomposer, env);
});
}
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds the largest K keys from an unordered input sequence of keys.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! - @devicestorage
//!
//! .. versionadded:: 3.3.0
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! The following code snippet demonstrates how to use the `cub::DeviceTopK::MinKeys` function to find the largest K
//! items:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-max-keys-non-deterministic-unsorted
//! :end-before: example-end topk-max-keys-non-deterministic-unsorted
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @param[in] d_temp_storage
//! @devicestorage
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in`
//!
//! @param[in] k
//! The value of K, which is the number of largest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, EnvT>, int> = 0>
CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceTopK::MaxKeys");
return detail::dispatch_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),
num_items,
k,
detail::identity_decomposer_t{},
::cuda::std::move(env));
}
//! @rst
//! Finds the largest K keys from an unordered input sequence.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! .. versionadded:: 3.5.0
//! First appears in CUDA Toolkit 13.5.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! Unlike the temp-storage overload, this overload allocates and manages the required temporary
//! storage internally using the memory resource queried from the environment.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-max-keys-env
//! :end-before: example-end topk-max-keys-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in`
//!
//! @param[in] k
//! The value of K, which is the number of largest keys to find from `num_items` keys. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, EnvT>, int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t
MaxKeys(KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceTopK::MaxKeys");
return detail::dispatch_with_env(
env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) {
return detail::dispatch_topk<detail::topk::select::max>(
storage,
bytes,
d_keys_in,
d_keys_out,
static_cast<NullType*>(nullptr),
static_cast<NullType*>(nullptr),
num_items,
k,
detail::identity_decomposer_t{},
env);
});
}
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds the largest K keys from an unordered input sequence of keys,
//! using a decomposer to interpret user-defined key types.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! - @devicestorage
//!
//! .. versionadded:: 3.4.0
//! First appears in CUDA Toolkit 13.4.
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Let's consider a user-defined ``custom_t`` type below. To find the top-k
//! elements of an array of ``custom_t`` objects, we have to tell CUB about
//! relevant members of the ``custom_t`` type. We do this by providing a
//! decomposer that returns a tuple of references to relevant members of the key.
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-custom-type
//! :end-before: example-end topk-custom-type
//!
//! The following snippet shows how to find the top-k largest keys of ``custom_t``
//! objects using ``cub::DeviceTopK::MaxKeys``:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-max-keys-custom-type
//! :end-before: example-end topk-max-keys-custom-type
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam DecomposerT
//! **[inferred]** Type of a callable object responsible for decomposing a key into a tuple of references to its
//! constituent arithmetic types.
//!
//! @param[in] d_temp_storage
//! @devicestorage
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in`
//!
//! @param[in] k
//! The value of K, which is the number of largest keys to find from `num_items` keys. Capped to a maximum of
//! `num_items`.
//!
//! @param decomposer
//! Callable object responsible for decomposing a key into a tuple of references to its constituent arithmetic
//! types.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`.
//! @endrst
template <typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT = ::cuda::std::execution::env<>>
CUB_RUNTIME_FUNCTION static //
::cuda::std::enable_if_t<detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, DecomposerT>,
cudaError_t>
MaxKeys(void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceTopK::MaxKeys");
using key_t = detail::it_value_t<KeyInputIteratorT>;
static_assert(!detail::radix::can_twiddle<key_t>,
"Custom decomposers are not supported for fundamental types; "
"use the non-decomposer API overload instead");
return detail::dispatch_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),
num_items,
k,
decomposer,
::cuda::std::move(env));
}
//! @rst
//! Finds the largest K keys from an unordered input sequence,
//! using a decomposer to interpret user-defined key types.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! .. versionadded:: 3.5.0
//! First appears in CUDA Toolkit 13.5.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! Unlike the temp-storage overload, this overload allocates and manages the required temporary
//! storage internally using the memory resource queried from the environment.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-max-keys-decomposer-env
//! :end-before: example-end topk-max-keys-decomposer-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam DecomposerT
//! **[inferred]** Type of a callable object responsible for decomposing a key into a tuple of references to its
//! constituent arithmetic types.
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in`
//!
//! @param[in] k
//! The value of K, which is the number of largest keys to find from `num_items` keys. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] decomposer
//! Callable object responsible for decomposing a key into a tuple of references to its constituent arithmetic
//! types.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, DecomposerT>,
int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MaxKeys(
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceTopK::MaxKeys");
using key_t = detail::it_value_t<KeyInputIteratorT>;
static_assert(!detail::radix::can_twiddle<key_t>,
"Custom decomposers are not supported for fundamental types; "
"use the non-decomposer API overload instead");
return detail::dispatch_with_env(
env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) {
return detail::dispatch_topk<detail::topk::select::max>(
storage,
bytes,
d_keys_in,
d_keys_out,
static_cast<NullType*>(nullptr),
static_cast<NullType*>(nullptr),
num_items,
k,
decomposer,
env);
});
}
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds the lowest K keys from an unordered input sequence of keys.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! - @devicestorage
//!
//! .. versionadded:: 3.3.0
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! The following code snippet demonstrates how to use the `cub::DeviceTopK::MinKeys` function to find the lowest K
//! items:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-min-keys-non-deterministic-unsorted
//! :end-before: example-end topk-min-keys-non-deterministic-unsorted
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @param[in] d_temp_storage
//! @devicestorage
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in`
//!
//! @param[in] k
//! The value of K, which is the number of largest pairs to find from `num_items` pairs. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, EnvT>, int> = 0>
CUB_RUNTIME_FUNCTION static cudaError_t MinKeys(
void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceTopK::MinKeys");
return detail::dispatch_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),
num_items,
k,
detail::identity_decomposer_t{},
::cuda::std::move(env));
}
//! @rst
//! Finds the smallest K keys from an unordered input sequence.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! .. versionadded:: 3.5.0
//! First appears in CUDA Toolkit 13.5.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! Unlike the temp-storage overload, this overload allocates and manages the required temporary
//! storage internally using the memory resource queried from the environment.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-min-keys-env
//! :end-before: example-end topk-min-keys-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in`
//!
//! @param[in] k
//! The value of K, which is the number of lowest keys to find from `num_items` keys. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<!detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, EnvT>, int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t
MinKeys(KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
NumItemsT num_items,
NumOutItemsT k,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceTopK::MinKeys");
return detail::dispatch_with_env(
env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) {
return detail::dispatch_topk<detail::topk::select::min>(
storage,
bytes,
d_keys_in,
d_keys_out,
static_cast<NullType*>(nullptr),
static_cast<NullType*>(nullptr),
num_items,
k,
detail::identity_decomposer_t{},
env);
});
}
//! @rst
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Finds the lowest K keys from an unordered input sequence of keys,
//! using a decomposer to interpret user-defined key types.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! - @devicestorage
//!
//! .. versionadded:: 3.4.0
//! First appears in CUDA Toolkit 13.4.
//!
//! A Simple Example
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! Let's consider a user-defined ``custom_t`` type below. To find the top-k
//! elements of an array of ``custom_t`` objects, we have to tell CUB about
//! relevant members of the ``custom_t`` type. We do this by providing a
//! decomposer that returns a tuple of references to relevant members of the key.
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-custom-type
//! :end-before: example-end topk-custom-type
//!
//! The following snippet shows how to find the top-k smallest keys of ``custom_t``
//! objects using ``cub::DeviceTopK::MinKeys``:
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-min-keys-custom-type
//! :end-before: example-end topk-min-keys-custom-type
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam DecomposerT
//! **[inferred]** Type of a callable object responsible for decomposing a key into a tuple of references to its
//! constituent arithmetic types.
//!
//! @param[in] d_temp_storage
//! @devicestorage
//!
//! @param[in,out] temp_storage_bytes
//! Reference to size in bytes of `d_temp_storage` allocation
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in`
//!
//! @param[in] k
//! The value of K, which is the number of lowest keys to find from `num_items` keys. Capped to a maximum of
//! `num_items`.
//!
//! @param decomposer
//! Callable object responsible for decomposing a key into a tuple of references to its constituent arithmetic
//! types.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`.
//! @endrst
template <typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT = ::cuda::std::execution::env<>>
CUB_RUNTIME_FUNCTION static //
::cuda::std::enable_if_t<detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, DecomposerT>,
cudaError_t>
MinKeys(void* d_temp_storage,
size_t& temp_storage_bytes,
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceTopK::MinKeys");
using key_t = detail::it_value_t<KeyInputIteratorT>;
static_assert(!detail::radix::can_twiddle<key_t>,
"Custom decomposers are not supported for fundamental types; "
"use the non-decomposer API overload instead");
return detail::dispatch_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),
num_items,
k,
decomposer,
::cuda::std::move(env));
}
//! @rst
//! Finds the smallest K keys from an unordered input sequence,
//! using a decomposer to interpret user-defined key types.
//!
//! .. note::
//!
//! The behavior is undefined if the input and output ranges overlap in any way.
//!
//! .. versionadded:: 3.5.0
//! First appears in CUDA Toolkit 13.5.
//!
//! This is an environment-based API that allows customization of:
//!
//! - Stream: Query via ``cuda::get_stream``
//! - Memory resource: Query via ``cuda::mr::get_memory_resource``
//!
//! Unlike the temp-storage overload, this overload allocates and manages the required temporary
//! storage internally using the memory resource queried from the environment.
//!
//! Snippet
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! .. literalinclude:: ../../../cub/test/catch2_test_device_topk_env_api.cu
//! :language: c++
//! :dedent:
//! :start-after: example-begin topk-min-keys-decomposer-env
//! :end-before: example-end topk-min-keys-decomposer-env
//!
//! @endrst
//!
//! @tparam KeyInputIteratorT
//! **[inferred]** Random-access input iterator type for reading input keys @iterator
//!
//! @tparam KeyOutputIteratorT
//! **[inferred]** Random-access output iterator type for writing output keys @iterator
//!
//! @tparam NumItemsT
//! The integral type of variable num_items
//!
//! @tparam NumOutItemsT
//! The integral type of variable k
//!
//! @tparam DecomposerT
//! **[inferred]** Type of a callable object responsible for decomposing a key into a tuple of references to its
//! constituent arithmetic types.
//!
//! @tparam EnvT
//! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``.
//!
//! @param[in] d_keys_in
//! Random-access iterator to the input sequence containing the keys
//!
//! @param[out] d_keys_out
//! Random-access iterator to the output sequence of keys, where K values will be written to
//!
//! @param[in] num_items
//! Number of items to be read and processed from `d_keys_in`
//!
//! @param[in] k
//! The value of K, which is the number of lowest keys to find from `num_items` keys. Capped to a maximum of
//! `num_items`.
//!
//! @param[in] decomposer
//! Callable object responsible for decomposing a key into a tuple of references to its constituent arithmetic
//! types.
//!
//! @param[in] env
//! @rst
//! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``.
//! @endrst
template <
typename KeyInputIteratorT,
typename KeyOutputIteratorT,
typename NumItemsT,
typename NumOutItemsT,
typename DecomposerT,
typename EnvT = ::cuda::std::execution::env<>,
::cuda::std::enable_if_t<detail::radix::is_valid_decomposer<detail::it_value_t<KeyInputIteratorT>, DecomposerT>,
int> = 0>
[[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MinKeys(
KeyInputIteratorT d_keys_in,
KeyOutputIteratorT d_keys_out,
NumItemsT num_items,
NumOutItemsT k,
DecomposerT decomposer,
const EnvT& env = {})
{
_CCCL_NVTX_RANGE_SCOPE("cub::DeviceTopK::MinKeys");
using key_t = detail::it_value_t<KeyInputIteratorT>;
static_assert(!detail::radix::can_twiddle<key_t>,
"Custom decomposers are not supported for fundamental types; "
"use the non-decomposer API overload instead");
return detail::dispatch_with_env(
env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, [[maybe_unused]] auto stream) {
return detail::dispatch_topk<detail::topk::select::min>(
storage,
bytes,
d_keys_in,
d_keys_out,
static_cast<NullType*>(nullptr),
static_cast<NullType*>(nullptr),
num_items,
k,
decomposer,
env);
});
}
};
CUB_NAMESPACE_END
Reference in New Issue View Git Blame Copy Permalink