diff --git a/qwen3_6_scripts/build_cccl_moe_sort_scatter.sh b/qwen3_6_scripts/build_cccl_moe_sort_scatter.sh new file mode 100644 index 00000000..0511b51c --- /dev/null +++ b/qwen3_6_scripts/build_cccl_moe_sort_scatter.sh @@ -0,0 +1,76 @@ +#!/usr/bin/env bash +# Build cccl_moe_sort_scatter.so using CCCL upstream headers +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SRC="${SCRIPT_DIR}/cccl_moe_sort_scatter.cu" +INC="${SCRIPT_DIR}/cccl_preload/include" +OUT="${1:-${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10/cccl_moe_sort_scatter.so}" + +[[ -f "${SRC}" ]] || { echo "Source not found: ${SRC}"; exit 2; } +[[ -d "${INC}/cub" ]] || { echo "CCCL include tree missing: ${INC}/cub"; exit 2; } + +# Find NVCC or corex clang +NVCC="" +for candidate in \ + /usr/local/corex/bin/nvcc \ + /usr/local/cuda/bin/nvcc \ + ; do + if [[ -x "${candidate}" ]]; then + NVCC="${candidate}" + break + fi +done + +TORCH_INC=$(python3 -c "import torch; print(torch.utils.cpp_extension.include_paths()[0])" 2>/dev/null) +TORCH_LIB=$(python3 -c "import torch; print(torch.utils.cmake_prefix_path + '/../lib')" 2>/dev/null || echo "") +PYTHON_INC=$(python3 -c "from sysconfig import get_paths; print(get_paths()['include'])" 2>/dev/null) + +echo "[build] NVCC: ${NVCC:-not found}" +echo "[build] CCCL: ${INC}" +echo "[build] Torch: ${TORCH_INC}" +echo "[build] Output: ${OUT}" + +if [[ -n "${NVCC}" ]]; then + "${NVCC}" \ + -shared --compiler-options -fPIC \ + -O3 -std=c++17 \ + -I"${INC}" \ + -I"${TORCH_INC}" \ + -I"${TORCH_INC}/torch/csrc/api/include" \ + ${PYTHON_INC:+-I"${PYTHON_INC}"} \ + -DCUB_WRAPPED_NAMESPACE=cccl_moe \ + -DTORCH_EXTENSION_NAME=cccl_moe_sort_scatter \ + -x cu \ + -o "${OUT}" "${SRC}" \ + -ltorch -lc10 -ltorch_cuda -ltorch_cpu \ + ${TORCH_LIB:+-L"${TORCH_LIB}"} \ + 2>&1 +else + echo "[build] No nvcc found, trying torch JIT at runtime" + python3 -c " +from torch.utils.cpp_extension import load +mod = load( + name='cccl_moe_sort_scatter', + sources=['${SRC}'], + extra_include_paths=['${INC}'], + extra_cuda_cflags=['-O3', '-DCUB_WRAPPED_NAMESPACE=cccl_moe'], + verbose=True, +) +print('[build] JIT compiled successfully') +import shutil, os +# Copy to output +src_so = os.path.join(os.path.dirname(mod.__file__), 'cccl_moe_sort_scatter.so') +if os.path.exists(src_so): + os.makedirs(os.path.dirname('${OUT}'), exist_ok=True) + shutil.copy2(src_so, '${OUT}') + print(f'[build] Copied to ${OUT}') +" 2>&1 +fi + +if [[ -f "${OUT}" ]]; then + echo "[build] SUCCESS: ${OUT} ($(stat -c%s "${OUT}" 2>/dev/null || echo '?') bytes)" +else + echo "[build] FAILED" + exit 1 +fi diff --git a/qwen3_6_scripts/cccl_moe_sort_scatter.cu b/qwen3_6_scripts/cccl_moe_sort_scatter.cu new file mode 100644 index 00000000..dd57d92a --- /dev/null +++ b/qwen3_6_scripts/cccl_moe_sort_scatter.cu @@ -0,0 +1,201 @@ +// cccl_moe_sort_scatter.cu — CUB DeviceRadixSort-based MoE token dispatch +// +// Replaces the 3-kernel (histogram + prefix_sum + place) approach with: +// 1. DeviceRadixSort::SortPairs — sort (expert_id, token_idx) pairs by expert_id +// 2. DeviceHistogram::HistogramEven — count tokens per expert +// 3. DeviceScan::ExclusiveSum — prefix sum for expert offsets +// +// Uses CCCL upstream headers (in cccl_preload/include/) instead of corex CUB +// to avoid BI-V100 corex CUB bugs. +// +// Build: torch.utils.cpp_extension.load( +// name="cccl_moe_sort_scatter", +// sources=["cccl_moe_sort_scatter.cu"], +// extra_include_paths=["cccl_preload/include"], +// extra_cuda_cflags=["-O3", "-DCUB_WRAPPED_NAMESPACE=cccl_moe"], +// ) + +#include +#include +#include + +// Use CCCL CUB, not corex CUB +#define CUB_WRAPPED_NAMESPACE cccl_moe +#include +#include +#include + +// ======================================================================== +// moe_sort_scatter: sort tokens by expert_id using CUB DeviceRadixSort +// +// Input: +// expert_id: [N] int32, each in [0, num_experts) +// num_experts: int +// +// Output: +// sorted_indices: [N] int32 — original token indices sorted by expert +// expert_offsets: [num_experts+1] int32 — exclusive prefix sum +// expert_sizes: [num_experts] int32 — count per expert +// ======================================================================== + +// Small kernel to build expert_sizes from sorted keys via boundary detection +__global__ void compute_expert_boundaries( + const int32_t* __restrict__ sorted_keys, + int32_t* __restrict__ expert_offsets, // [num_experts + 1] + int64_t N, + int32_t num_experts) { + // Initialize all to 0 + int tid = blockIdx.x * blockDim.x + threadIdx.x; + + // First pass: detect boundaries + if (tid < N) { + int32_t cur = sorted_keys[tid]; + if (tid == 0) { + // First element starts expert cur + expert_offsets[cur] = 0; + } else { + int32_t prev = sorted_keys[tid - 1]; + if (cur != prev) { + expert_offsets[cur] = tid; + } + } + // Last element + if (tid == N - 1) { + expert_offsets[num_experts] = N; + } + } +} + +// Fill gaps in expert_offsets (experts with 0 tokens) +__global__ void fill_offset_gaps( + int32_t* __restrict__ expert_offsets, + int32_t num_experts) { + // Backward fill: if expert_offsets[i] == -1, copy from next non-(-1) + // Single thread is fine for num_experts <= 256 + if (threadIdx.x != 0) return; + + // Fill from the end + int32_t next_offset = expert_offsets[num_experts]; // = N + for (int i = num_experts - 1; i >= 0; --i) { + if (expert_offsets[i] == -1) { + expert_offsets[i] = next_offset; + } else { + next_offset = expert_offsets[i]; + } + } +} + +// Compute expert_sizes from expert_offsets +__global__ void compute_expert_sizes( + const int32_t* __restrict__ expert_offsets, + int32_t* __restrict__ expert_sizes, + int32_t num_experts) { + int tid = blockIdx.x * blockDim.x + threadIdx.x; + if (tid < num_experts) { + expert_sizes[tid] = expert_offsets[tid + 1] - expert_offsets[tid]; + } +} + + +std::tuple +moe_sort_scatter(const torch::Tensor& expert_id, int64_t num_experts) { + TORCH_CHECK(expert_id.is_cuda(), "expert_id must be on CUDA"); + auto device = expert_id.device(); + auto stream = at::cuda::getCurrentCUDAStream(); + int64_t N = expert_id.numel(); + int32_t E = static_cast(num_experts); + + // Ensure int32 + auto keys_in = expert_id.to(torch::kInt32).contiguous(); + auto opt_i32 = keys_in.options(); + + // Create value array: [0, 1, 2, ..., N-1] + auto vals_in = torch::arange(N, opt_i32); + + // Allocate output + auto sorted_keys = torch::empty({N}, opt_i32); + auto sorted_vals = torch::empty({N}, opt_i32); + + // CUB DeviceRadixSort needs temp storage + // First query size + size_t temp_bytes = 0; + cccl_moe::cub::DeviceRadixSort::SortPairs( + nullptr, temp_bytes, + keys_in.data_ptr(), + sorted_keys.data_ptr(), + vals_in.data_ptr(), + sorted_vals.data_ptr(), + static_cast(N), + 0, // begin_bit + sizeof(int32_t) * 8, // end_bit (all bits, but only need log2(E) bits) + stream); + + // Allocate temp storage + auto temp_storage = torch::empty({static_cast(temp_bytes)}, + torch::dtype(torch::kUInt8).device(device)); + + // Sort + cccl_moe::cub::DeviceRadixSort::SortPairs( + temp_storage.data_ptr(), temp_bytes, + keys_in.data_ptr(), + sorted_keys.data_ptr(), + vals_in.data_ptr(), + sorted_vals.data_ptr(), + static_cast(N), + 0, + sizeof(int32_t) * 8, + stream); + + // Compute expert offsets via boundary detection + // Initialize to -1 + auto expert_offsets = torch::full({num_experts + 1}, -1, opt_i32); + + int block = 256; + int grid = (N + block - 1) / block; + compute_expert_boundaries<<>>( + sorted_keys.data_ptr(), + expert_offsets.data_ptr(), + N, E); + + // Handle empty experts + fill_offset_gaps<<<1, 1, 0, stream>>>( + expert_offsets.data_ptr(), E); + + // Compute sizes from offsets + auto expert_sizes = torch::empty({num_experts}, opt_i32); + int grid2 = (E + block - 1) / block; + compute_expert_sizes<<>>( + expert_offsets.data_ptr(), + expert_sizes.data_ptr(), + E); + + // sorted_vals = the original token indices, sorted by expert + // expert_sizes = tokens per expert + // sorted_keys not needed by caller, but sorted_vals is "dst_src" + // + // Build src_dst: inverse mapping + // src_dst[sorted_vals[i]] = i + auto src_dst = torch::empty({N}, opt_i32); + + // Simple inverse scatter kernel + // For now use a tiny lambda — could be another kernel + // But actually we can do it with scatter: + // src_dst.scatter_(0, sorted_vals.long(), torch.arange(N)) + // This is a single CUDA kernel internally + + auto arange_n = torch::arange(N, opt_i32); + src_dst.scatter_(0, sorted_vals.to(torch::kInt64), arange_n); + + return std::make_tuple(src_dst, sorted_vals, expert_sizes); +} + + +// ======================================================================== +// pybind +// ======================================================================== + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("moe_sort_scatter", &moe_sort_scatter, + "CUB DeviceRadixSort-based MoE token dispatch " + "(sort expert_ids, compute offsets+sizes)"); +} diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_adjacent_difference.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_adjacent_difference.cuh new file mode 100644 index 00000000..4f142c7c --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_adjacent_difference.cuh @@ -0,0 +1,245 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +template +struct agent_adjacent_difference_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int ITEMS_PER_THREAD = ItemsPerThread; + static constexpr int ITEMS_PER_TILE = BLOCK_THREADS * ITEMS_PER_THREAD; + + static constexpr cub::BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + static constexpr cub::CacheLoadModifier LOAD_MODIFIER = LoadModifier; + static constexpr cub::BlockStoreAlgorithm STORE_ALGORITHM = StoreAlgorithm; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template +using AgentAdjacentDifferencePolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceAdjacentDifference") = + detail::agent_adjacent_difference_policy; + +namespace detail::adjacent_difference +{ +template +struct AgentDifference +{ + using LoadIt = try_make_cache_modified_iterator_t; + + using BlockLoad = typename cub::BlockLoadType::type; + using BlockStore = typename cub::BlockStoreType::type; + + using BlockAdjacentDifferenceT = cub::BlockAdjacentDifference; + + union _TempStorage + { + typename BlockLoad::TempStorage load; + typename BlockStore::TempStorage store; + typename BlockAdjacentDifferenceT::TempStorage adjacent_difference; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + static constexpr int BLOCK_THREADS = Policy::BLOCK_THREADS; + static constexpr int ITEMS_PER_THREAD = Policy::ITEMS_PER_THREAD; + static constexpr int ITEMS_PER_TILE = Policy::ITEMS_PER_TILE; + static constexpr int SHARED_MEMORY_SIZE = static_cast(sizeof(TempStorage)); + + _TempStorage& temp_storage; + InputIteratorT input_it; + LoadIt load_it; + InputT* first_tile_previous; + OutputIteratorT result; + DifferenceOpT difference_op; + OffsetT num_items; + + _CCCL_DEVICE _CCCL_FORCEINLINE AgentDifference( + TempStorage& temp_storage, + InputIteratorT input_it, + InputT* first_tile_previous, + OutputIteratorT result, + DifferenceOpT difference_op, + OffsetT num_items) + : temp_storage(temp_storage.Alias()) + , input_it(input_it) + , load_it(try_make_cache_modified_iterator(input_it)) + , first_tile_previous(first_tile_previous) + , result(result) + , difference_op(difference_op) + , num_items(num_items) + {} + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void consume_tile_impl(int num_remaining, int tile_idx, OffsetT tile_base) + { + InputT input[ITEMS_PER_THREAD]; + OutputT output[ITEMS_PER_THREAD]; + + if (IS_LAST_TILE) + { + // Fill last elements with the first element + // because collectives are not suffix guarded + BlockLoad(temp_storage.load).Load(load_it + tile_base, input, num_remaining, *(load_it + tile_base)); + } + else + { + BlockLoad(temp_storage.load).Load(load_it + tile_base, input); + } + + __syncthreads(); + + if (ReadLeft) + { + if (IS_FIRST_TILE) + { + if (IS_LAST_TILE) + { + BlockAdjacentDifferenceT(temp_storage.adjacent_difference) + .SubtractLeftPartialTile(input, output, difference_op, num_remaining); + } + else + { + BlockAdjacentDifferenceT(temp_storage.adjacent_difference).SubtractLeft(input, output, difference_op); + } + } + else + { + InputT tile_prev_input = MayAlias ? first_tile_previous[tile_idx] : *(input_it + tile_base - 1); + + if (IS_LAST_TILE) + { + BlockAdjacentDifferenceT(temp_storage.adjacent_difference) + .SubtractLeftPartialTile(input, output, difference_op, num_remaining, tile_prev_input); + } + else + { + BlockAdjacentDifferenceT(temp_storage.adjacent_difference) + .SubtractLeft(input, output, difference_op, tile_prev_input); + } + } + } + else + { + if (IS_LAST_TILE) + { + BlockAdjacentDifferenceT(temp_storage.adjacent_difference) + .SubtractRightPartialTile(input, output, difference_op, num_remaining); + } + else + { + InputT tile_next_input = MayAlias ? first_tile_previous[tile_idx] : *(input_it + tile_base + ITEMS_PER_TILE); + + BlockAdjacentDifferenceT(temp_storage.adjacent_difference) + .SubtractRight(input, output, difference_op, tile_next_input); + } + } + + __syncthreads(); + + if (IS_LAST_TILE) + { + BlockStore(temp_storage.store).Store(result + tile_base, output, num_remaining); + } + else + { + BlockStore(temp_storage.store).Store(result + tile_base, output); + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void consume_tile(int num_remaining, int tile_idx, OffsetT tile_base) + { + if (tile_idx == 0) + { + consume_tile_impl(num_remaining, tile_idx, tile_base); + } + else + { + consume_tile_impl(num_remaining, tile_idx, tile_base); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void Process(int tile_idx, OffsetT tile_base) + { + OffsetT num_remaining = num_items - tile_base; + + if (num_remaining > ITEMS_PER_TILE) // not a last tile + { + consume_tile(num_remaining, tile_idx, tile_base); + } + else + { + consume_tile(num_remaining, tile_idx, tile_base); + } + } +}; + +template +struct AgentDifferenceInit +{ + static constexpr int BLOCK_THREADS = 128; + + static _CCCL_DEVICE _CCCL_FORCEINLINE void + Process(int tile_idx, InputIteratorT first, InputT* result, OffsetT num_tiles, int items_per_tile) + { + OffsetT tile_base = static_cast(tile_idx) * items_per_tile; + + if (tile_base > 0 && tile_idx < num_tiles) + { + if (ReadLeft) + { + result[tile_idx] = first[tile_base - 1]; + } + else + { + result[tile_idx - 1] = first[tile_base]; + } + } + } +}; +} // namespace detail::adjacent_difference + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_batch_memcpy.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_batch_memcpy.cuh new file mode 100644 index 00000000..093ea61b --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_batch_memcpy.cuh @@ -0,0 +1,1164 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * cub::AgentBatchMemcpy implements device-wide copying of a batch of device-accessible + * source-buffers to device-accessible destination-buffers. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +// TODO(bgruber): drop in CCCL 4.0 +template +struct agent_batch_memcpy_policy +{ + static constexpr uint32_t BLOCK_THREADS = ThreadsPerBlock; + static constexpr uint32_t BUFFERS_PER_THREAD = BuffersPerThread; + static constexpr uint32_t TLEV_BYTES_PER_THREAD = TlevBytesPerThread; + static constexpr uint32_t PREFER_POW2_BITS = PreferPow2Bits; + static constexpr uint32_t BLOCK_LEVEL_TILE_SIZE = BlockLevelTileSize; + static constexpr uint32_t WARP_LEVEL_THRESHOLD = WarpLevelThreshold; + static constexpr uint32_t BLOCK_LEVEL_THRESHOLD = BlockLevelThreshold; + + using buff_delay_constructor = BuffDelayConstructor; + using block_delay_constructor = BlockDelayConstructor; +}; +} // namespace detail + +// TODO(bgruber): drop in CCCL 4.0 +//! Deprecated [Since 3.5] +template +using AgentBatchMemcpyPolicy + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceMemcpy") = detail::agent_batch_memcpy_policy< + ThreadsPerBlock, + BuffersPerThread, + TlevBytesPerThread, + PreferPow2Bits, + BlockLevelTileSize, + WarpLevelThreshold, + BlockLevelThreshold, + BuffDelayConstructor, + BlockDelayConstructor>; + +namespace detail::batch_memcpy +{ +template +_CCCL_FORCEINLINE _CCCL_DEVICE void +LoadVectorAndFunnelShiftR(uint32_t const* aligned_ptr, uint32_t bit_shift, uint4& data_out) +{ + data_out = {aligned_ptr[0], aligned_ptr[1], aligned_ptr[2], aligned_ptr[3]}; + + if constexpr (!PTR_IS_FOUR_BYTE_ALIGNED) + { + uint32_t tail = aligned_ptr[4]; + data_out.x = __funnelshift_r(data_out.x, data_out.y, bit_shift); + data_out.y = __funnelshift_r(data_out.y, data_out.z, bit_shift); + data_out.z = __funnelshift_r(data_out.z, data_out.w, bit_shift); + data_out.w = __funnelshift_r(data_out.w, tail, bit_shift); + } +} + +template +_CCCL_FORCEINLINE _CCCL_DEVICE void +LoadVectorAndFunnelShiftR(uint32_t const* aligned_ptr, uint32_t bit_shift, uint2& data_out) +{ + data_out = {aligned_ptr[0], aligned_ptr[1]}; + + if constexpr (!PTR_IS_FOUR_BYTE_ALIGNED) + { + uint32_t tail = aligned_ptr[2]; + data_out.x = __funnelshift_r(data_out.x, data_out.y, bit_shift); + data_out.y = __funnelshift_r(data_out.y, tail, bit_shift); + } +} + +template +_CCCL_FORCEINLINE _CCCL_DEVICE void +LoadVectorAndFunnelShiftR(uint32_t const* aligned_ptr, uint32_t bit_shift, uint32_t& data_out) +{ + data_out = aligned_ptr[0]; + + if constexpr (!PTR_IS_FOUR_BYTE_ALIGNED) + { + uint32_t tail = aligned_ptr[1]; + data_out = __funnelshift_r(data_out, tail, bit_shift); + } +} + +/** + * @brief Loads data from \p ptr into \p data_out without requiring \p ptr to be aligned. + * @note If \p ptr isn't aligned to four bytes, the bytes from the last four-byte aligned address up + * to \p ptr are loaded too (but dropped) and, hence, need to be device-accessible. Similarly, if + * \p ptr isn't aligned to four bytes, the bytes from `(ptr + sizeof(VectorT))` up to the following + * four-byte aligned address are loaded too (but dropped), and, hence, need to be device-accessible. + * + * @tparam VectorT The vector type used for vectorized stores (i.e., one of uint4, uint2, uint32_t) + * @param ptr The pointer from which the data is supposed to be loaded + * @param data_out The vector type that stores the data loaded from \p ptr + */ +template +_CCCL_FORCEINLINE _CCCL_DEVICE void LoadVector(const char* ptr, VectorT& data_out) +{ + const uint32_t offset = reinterpret_cast(ptr) % 4U; + const uint32_t* aligned_ptr = reinterpret_cast(ptr - offset); + constexpr uint32_t bits_per_byte = 8U; + const uint32_t bit_shift = offset * bits_per_byte; + + // If `ptr` is aligned to four bytes, we can perform a simple uint32_t-aliased load + if (offset == 0) + { + LoadVectorAndFunnelShiftR(aligned_ptr, bit_shift, data_out); + } + // Otherwise, we need to load extra bytes and perform funnel-shifting + else + { + LoadVectorAndFunnelShiftR(aligned_ptr, bit_shift, data_out); + } +} + +/** + * @brief Helper data structure to hold information on the byte range for which we can safely + * perform vectorized copies. + * + * @tparam VectorT The vector type used for vectorized stores (i.e., one of uint4, uint2, uint32_t) + */ +template +struct PointerRange +{ + VectorT* out_begin; + VectorT* out_end; + const char* in_begin; + const char* in_end; +}; + +/** + * @brief Both `out_start_aligned` and `out_end_aligned` are indices into `out_ptr`. + * `out_start_aligned` is the first VectorT-aligned memory location after `out_ptr + 3`. + * `out_end_aligned` is the last VectorT-aligned memory location before `out_end - 4`, where out_end + * corresponds to one past the last byte to be copied. Bytes between `[out_start_aligned, + * out_end_aligned)` will be copied using VectorT. `out_ptr + 3` and `out_end - 4` are used instead + * of `out_ptr` and `out_end` to avoid `LoadVector` reading beyond data boundaries. + * + * @tparam VectorT The vector type used for vectorized stores (i.e., one of uint4, uint2, uint32_t) + * @tparam ByteOffsetT Type used to index the bytes within the buffers + * @param in_begin Pointer to the beginning of the byte range that shall be copied + * @param out_begin Pointer to the beginning of the byte range that shall be copied + * @param num_bytes Number of bytes that shall be copied + * @return The byte range that can safely be copied using vectorized stores of type VectorT + */ +template +_CCCL_DEVICE _CCCL_FORCEINLINE PointerRange +GetAlignedPtrs(const void* in_begin, void* out_begin, ByteOffsetT num_bytes) +{ + // Data type size used for vectorized stores + constexpr auto out_datatype_size = uint32_t{sizeof(VectorT)}; + // Data type size used for type-aliased loads + constexpr auto in_datatype_size = uint32_t{sizeof(uint32_t)}; + + // char-aliased ptrs to simplify pointer arithmetic + char* out_ptr = static_cast(out_begin); + const char* in_ptr = static_cast(in_begin); + + // Number of bytes between the first VectorT-aligned address at or before out_begin and out_begin + const uint32_t alignment_offset = reinterpret_cast(out_ptr) % out_datatype_size; + + // The first VectorT-aligned address before (or at) out_begin + char* out_chars_aligned = out_ptr - alignment_offset; + + // The number of extra bytes preceding `in_ptr` that are loaded but dropped + uint32_t in_extra_bytes = reinterpret_cast(in_ptr) % in_datatype_size; + + // The offset required by `LoadVector`: + // If the input pointer is not aligned, we load data from the last aligned address preceding the + // pointer. That is, loading up to (in_datatype_size-1) bytes before `in_ptr` + uint32_t in_offset_req = in_extra_bytes; + + // Bytes after `out_chars_aligned` to the first VectorT-aligned address at or after `out_begin` + uint32_t out_start_aligned = ::cuda::round_up(in_offset_req + alignment_offset, out_datatype_size); + + // Compute the beginning of the aligned ranges (output and input pointers) + VectorT* out_aligned_begin = reinterpret_cast(out_chars_aligned + out_start_aligned); + const char* in_aligned_begin = in_ptr + (reinterpret_cast(out_aligned_begin) - out_ptr); + + // If the aligned range is not aligned for the input pointer, we load up to (in_datatype_size-1) + // bytes after the last byte that is copied. That is, we always load four bytes up to the next + // aligned input address at a time. E.g., if the last byte loaded is one byte past the last + // aligned address we'll also load the three bytes after that byte. + uint32_t in_extra_bytes_from_aligned = (reinterpret_cast(in_aligned_begin) % in_datatype_size); + uint32_t in_end_padding_req = (in_datatype_size - in_extra_bytes_from_aligned) % in_datatype_size; + + // Bytes after `out_chars_aligned` to the last VectorT-aligned + // address at (or before) `out_begin` + `num_bytes` + uint32_t out_end_aligned{}; + if (in_end_padding_req + alignment_offset > num_bytes) // NOLINT(bugprone-misplaced-widening-cast) + { + out_end_aligned = out_start_aligned; + } + else + { + out_end_aligned = (num_bytes - in_end_padding_req + alignment_offset) / out_datatype_size * out_datatype_size; + } + + VectorT* out_aligned_end = reinterpret_cast(out_chars_aligned + out_end_aligned); + const char* in_aligned_end = in_ptr + (reinterpret_cast(out_aligned_end) - out_ptr); + + return {out_aligned_begin, out_aligned_end, in_aligned_begin, in_aligned_end}; +} + +/** + * @brief Cooperatively copies \p num_bytes from \p src to \p dest using vectorized stores of type + * \p VectorT for addresses within [dest, dest + num_bytes) that are aligned to \p VectorT. A + * byte-wise copy is used for byte-ranges that are not aligned to \p VectorT. + * + * @tparam LOGICAL_WARP_SIZE The number of threads cooperaing to copy the data; all threads within + * [0, `LOGICAL_WARP_SIZE`) must invoke this method with the same arguments + * @tparam VectorT The vector type used for vectorized stores (i.e., one of uint4, uint2, uint32_t) + * @tparam ByteOffsetT Type used to index the bytes within the buffers + * @param thread_rank The thread rank within the group that cooperates to copy the data must be + * within [0, `LOGICAL_WARP_SIZE`) + * @param dest Pointer to the memory location to copy to + * @param num_bytes Number of bytes to copy + * @param src Pointer to the memory location to copy from + */ +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +vectorized_copy(int32_t thread_rank, void* dest, ByteOffsetT num_bytes, const void* src) +{ + char* out_ptr = static_cast(dest); + const char* in_ptr = static_cast(src); + + // Gets the byte range that can safely be copied using vectorized stores of type VectorT + auto aligned_range = GetAlignedPtrs(src, dest, num_bytes); + + // If byte range for which we can use vectorized copies is empty -> use byte-wise copies + if (aligned_range.out_end <= aligned_range.out_begin) + { + for (ByteOffsetT ichar = thread_rank; ichar < num_bytes; ichar += LOGICAL_WARP_SIZE) + { + out_ptr[ichar] = in_ptr[ichar]; + } + } + else + { + // Copy bytes in range `[dest, aligned_range.out_begin)` + out_ptr += thread_rank; + in_ptr += thread_rank; + while (out_ptr < reinterpret_cast(aligned_range.out_begin)) + { + *out_ptr = *in_ptr; + out_ptr += LOGICAL_WARP_SIZE; + in_ptr += LOGICAL_WARP_SIZE; + } + + // Copy bytes in range `[aligned_range.out_begin, aligned_range.out_end)` + VectorT* aligned_range_begin = aligned_range.out_begin + thread_rank; + const char* in_aligned_begin = aligned_range.in_begin + thread_rank * sizeof(VectorT); + while (aligned_range_begin < aligned_range.out_end) + { + VectorT data_in; + LoadVector(in_aligned_begin, data_in); + *aligned_range_begin = data_in; + in_aligned_begin += sizeof(VectorT) * LOGICAL_WARP_SIZE; + aligned_range_begin += LOGICAL_WARP_SIZE; + } + + // Copy bytes in range `[aligned_range.out_end, dest + num_bytes)`. + out_ptr = reinterpret_cast(aligned_range.out_end) + thread_rank; + in_ptr = aligned_range.in_end + thread_rank; + while (out_ptr < static_cast(dest) + num_bytes) + { + *out_ptr = *in_ptr; + out_ptr += LOGICAL_WARP_SIZE; + in_ptr += LOGICAL_WARP_SIZE; + } + } +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +copy_items(InputBufferT input_buffer, OutputBufferT output_buffer, OffsetT num_items, OffsetT offset = 0) +{ + if constexpr (IsMemcpy) + { + vectorized_copy( + threadIdx.x % LOGICAL_WARP_SIZE, + &reinterpret_cast(output_buffer)[offset], + num_items, + &reinterpret_cast(input_buffer)[offset]); + } + else + { + output_buffer += offset; + input_buffer += offset; + for (OffsetT i = threadIdx.x % LOGICAL_WARP_SIZE; i < num_items; i += LOGICAL_WARP_SIZE) + { + *(output_buffer + i) = *(input_buffer + i); + } + } +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE AliasT read_item(InputIt buffer_src, OffsetT offset) +{ + if constexpr (IsMemcpy) + { + return *(reinterpret_cast(buffer_src) + offset); + } + else + { + return *(buffer_src + offset); + } +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void write_item(OutputIt buffer_dst, OffsetT offset, AliasT value) +{ + if constexpr (IsMemcpy) + { + *(reinterpret_cast(buffer_dst) + offset) = value; + } + else + { + *(buffer_dst + offset) = value; + } +} + +enum class prefer_power_of_two_bits_option +{ + no, + yes +}; + +/** + * @brief A helper class that allows threads to maintain multiple counters, where the counter that + * shall be incremented can be addressed dynamically without incurring register spillage. + * + * @tparam NumItems The number of counters to allocate + * @tparam MaxItemValue The maximum count that must be supported. + * @tparam PreferPowerOfTwoBits Whether the number of bits to dedicate to each counter should be a + * power-of-two. If enabled, this allows replacing integer multiplication with a bit-shift in + * exchange for higher register pressure. + * @tparam BackingUnitT The data type that is used to provide the bits of all the counters that + * shall be allocated. + */ +template +class bit_packed_counter +{ +private: + /// The minimum number of bits required to represent all values from [0, MaxItemValue] + static constexpr uint32_t MIN_BITS_PER_ITEM = + (MaxItemValue == 0U) ? 1U : cub::Log2(MaxItemValue + 1U)>::VALUE; + + /// The number of bits allocated for each item. For pre-Volta, we prefer a power-of-2 here to + /// have the compiler replace costly integer multiplication with bit-shifting. + static constexpr uint32_t BITS_PER_ITEM = + (PreferPowerOfTwoBits == prefer_power_of_two_bits_option::yes) + ? (0x01ULL << (cub::Log2(MIN_BITS_PER_ITEM)>::VALUE)) + : MIN_BITS_PER_ITEM; + + /// The number of bits that each backing data type can store + static constexpr uint32_t NUM_BITS_PER_UNIT = sizeof(BackingUnitT) * 8; + + /// The number of items that each backing data type can store + static constexpr uint32_t ITEMS_PER_UNIT = NUM_BITS_PER_UNIT / BITS_PER_ITEM; + + /// The number of bits the backing data type is actually making use of + static constexpr uint32_t USED_BITS_PER_UNIT = ITEMS_PER_UNIT * BITS_PER_ITEM; + + /// The number of backing data types required to store the given number of items + static constexpr uint32_t NUM_TOTAL_UNITS = ::cuda::ceil_div(NumItems, ITEMS_PER_UNIT); + + /// This is the net number of bit-storage provided by each unit (remainder bits are unused) + static constexpr uint32_t UNIT_MASK = + (USED_BITS_PER_UNIT >= (8U * sizeof(uint32_t))) ? 0xFFFFFFFF : (0x01U << USED_BITS_PER_UNIT) - 1; + /// This is the bit-mask for each item + static constexpr uint32_t ITEM_MASK = + (BITS_PER_ITEM >= (8U * sizeof(uint32_t))) ? 0xFFFFFFFF : (0x01U << BITS_PER_ITEM) - 1; + + //------------------------------------------------------------------------------ + // ACCESSORS + //------------------------------------------------------------------------------ + +public: + _CCCL_DEVICE _CCCL_FORCEINLINE uint32_t get(uint32_t index) const + { + const uint32_t target_offset = index * BITS_PER_ITEM; + uint32_t val = 0; + + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < NUM_TOTAL_UNITS; ++i) + { + // In case the bit-offset of the counter at is larger than the bit range of the + // current unit, the bit_shift amount will be larger than the bits provided by this unit. As + // C++'s bit-shift has undefined behaviour if the bits being shifted exceed the operand width, + // we use the PTX instruction `shr` to make sure behaviour is well-defined. + // Negative bit-shift amounts wrap around in unsigned integer math and are ultimately clamped. + const uint32_t bit_shift = target_offset - i * USED_BITS_PER_UNIT; + val |= detail::LogicShiftRight(data[i], bit_shift) & ITEM_MASK; + } + return val; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void add(uint32_t index, uint32_t value) + { + const uint32_t target_offset = index * BITS_PER_ITEM; + + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < NUM_TOTAL_UNITS; ++i) + { + // In case the bit-offset of the counter at is larger than the bit range of the + // current unit, the bit_shift amount will be larger than the bits provided by this unit. As + // C++'s bit-shift has undefined behaviour if the bits being shifted exceed the operand width, + // we use the PTX instruction `shl` to make sure behaviour is well-defined. + // Negative bit-shift amounts wrap around in unsigned integer math and are ultimately clamped. + const uint32_t bit_shift = target_offset - i * USED_BITS_PER_UNIT; + data[i] += detail::LogicShiftLeft(value, bit_shift) & UNIT_MASK; + } + } + + _CCCL_DEVICE bit_packed_counter operator+(const bit_packed_counter& rhs) const + { + bit_packed_counter result; + + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < NUM_TOTAL_UNITS; ++i) + { + result.data[i] = data[i] + rhs.data[i]; + } + return result; + } + + //------------------------------------------------------------------------------ + // MEMBER VARIABLES + //------------------------------------------------------------------------------ + +private: + BackingUnitT data[NUM_TOTAL_UNITS] = {}; +}; + +template +class AgentBatchMemcpy +{ +private: + //--------------------------------------------------------------------- + // CONFIGS / CONSTANTS + //--------------------------------------------------------------------- + // Tuning policy-based configurations + static constexpr uint32_t BLOCK_THREADS = AgentMemcpySmallBuffersPolicyT::BLOCK_THREADS; + static constexpr uint32_t BUFFERS_PER_THREAD = AgentMemcpySmallBuffersPolicyT::BUFFERS_PER_THREAD; + static constexpr uint32_t TLEV_BYTES_PER_THREAD = AgentMemcpySmallBuffersPolicyT::TLEV_BYTES_PER_THREAD; + static constexpr prefer_power_of_two_bits_option PREFER_POW2_BITS = + (AgentMemcpySmallBuffersPolicyT::PREFER_POW2_BITS) + ? prefer_power_of_two_bits_option::yes + : prefer_power_of_two_bits_option::no; + static constexpr uint32_t BLOCK_LEVEL_TILE_SIZE = AgentMemcpySmallBuffersPolicyT::BLOCK_LEVEL_TILE_SIZE; + + // Derived configs + static constexpr uint32_t BUFFERS_PER_BLOCK = BUFFERS_PER_THREAD * BLOCK_THREADS; + static constexpr uint32_t TLEV_BUFFERS_PER_THREAD = BUFFERS_PER_THREAD; + static constexpr uint32_t BLEV_BUFFERS_PER_THREAD = BUFFERS_PER_THREAD; + + static constexpr uint32_t WARP_LEVEL_THRESHOLD = AgentMemcpySmallBuffersPolicyT::WARP_LEVEL_THRESHOLD; + + static constexpr uint32_t BLOCK_LEVEL_THRESHOLD = AgentMemcpySmallBuffersPolicyT::BLOCK_LEVEL_THRESHOLD; + + static constexpr uint32_t BUFFER_STABLE_PARTITION = false; + + // Constants + enum : uint32_t + { + TLEV_SIZE_CLASS = 0, + WLEV_SIZE_CLASS, + BLEV_SIZE_CLASS, + NUM_SIZE_CLASSES, + }; + + //--------------------------------------------------------------------- + // TYPE DECLARATIONS + //--------------------------------------------------------------------- + /// Internal load/store type. For byte-wise memcpy, a single-byte type + using AliasT = typename ::cuda::std:: + conditional_t, lazy_trait>>::type; + + /// Types of the input and output buffers + using InputBufferT = it_value_t; + using OutputBufferT = it_value_t; + + /// Type that has to be sufficiently large to hold any of the buffers' sizes. + /// The BufferSizeIteratorT's value type must be convertible to this type. + using BufferSizeT = it_value_t; + + /// Type used to index into the tile of buffers that this thread block is assigned to. + using BlockBufferOffsetT = uint16_t; + + /// Internal type used to index into the bytes of and represent size of a TLEV buffer + using TLevBufferSizeT = uint16_t; + + /** + * @brief Helper struct to simplify BlockExchange within a single four-byte word + */ + struct ZippedTLevByteAssignment + { + // The buffer id within this tile + BlockBufferOffsetT tile_buffer_id; + + // Byte-offset within that buffer + TLevBufferSizeT buffer_byte_offset; + }; + + /** + * POD to keep track of pairs after having partitioned this tile's + * buffers by their size. + */ + struct BufferTuple + { + // Size is only valid (and relevant) for buffers that are use thread-level collaboration + TLevBufferSizeT size; + + // The buffer id relative to this tile (i.e., the buffer id within this tile) + BlockBufferOffsetT buffer_id; + }; + + // Load buffers in a striped arrangement if we do not want to perform a stable partitioning into + // small, medium, and large buffers, otherwise load them in a blocked arrangement + using BufferLoadT = + BlockLoad(BLOCK_THREADS), + static_cast(BUFFERS_PER_THREAD), + BUFFER_STABLE_PARTITION ? BLOCK_LOAD_WARP_TRANSPOSE : BLOCK_LOAD_STRIPED>; + + // A vectorized counter that will count the number of buffers that fall into each of the + // size-classes. Where the size class represents the collaboration level that is required to + // process a buffer. The collaboration level being either: + //-> (1) TLEV (thread-level collaboration), requiring one or multiple threads but not a FULL warp + // to collaborate + //-> (2) WLEV (warp-level collaboration), requiring a full warp to collaborate on a buffer + //-> (3) BLEV (block-level collaboration), requiring one or multiple thread blocks to collaborate + // on a buffer */ + using VectorizedSizeClassCounterT = bit_packed_counter; + + // Block-level scan used to compute the write offsets + using BlockSizeClassScanT = cub::BlockScan(BLOCK_THREADS)>; + + // + using BlockBLevTileCountScanT = cub::BlockScan(BLOCK_THREADS)>; + + // Block-level run-length decode algorithm to evenly distribute work of all buffers requiring + // thread-level collaboration + using BlockRunLengthDecodeT = + cub::BlockRunLengthDecode(BLOCK_THREADS), + static_cast(TLEV_BUFFERS_PER_THREAD), + static_cast(TLEV_BYTES_PER_THREAD)>; + + using BlockExchangeTLevT = + cub::BlockExchange(BLOCK_THREADS), + static_cast(TLEV_BYTES_PER_THREAD)>; + + using BLevBuffScanPrefixCallbackOpT = + TilePrefixCallbackOp, + BLevBufferOffsetTileState, + typename AgentMemcpySmallBuffersPolicyT::buff_delay_constructor>; + + using BLevBlockScanPrefixCallbackOpT = + TilePrefixCallbackOp, + BLevBlockOffsetTileState, + typename AgentMemcpySmallBuffersPolicyT::block_delay_constructor>; + + //----------------------------------------------------------------------------- + // SHARED MEMORY DECLARATIONS + //----------------------------------------------------------------------------- + struct _TempStorage + { + union + { + typename BufferLoadT::TempStorage load_storage; + + // Stage 1: histogram over the size classes in preparation for partitioning buffers by size + typename BlockSizeClassScanT::TempStorage size_scan_storage; + + // Stage 2: Communicate the number ofer buffers requiring block-level collaboration + typename BLevBuffScanPrefixCallbackOpT::TempStorage buffer_scan_callback; + + // Stage 3; batch memcpy buffers that require only thread-level collaboration + struct + { + BufferTuple buffers_by_size_class[BUFFERS_PER_BLOCK]; + + // Stage 3.1: Write buffers requiring block-level collaboration to queue + union + { + struct + { + typename BLevBlockScanPrefixCallbackOpT::TempStorage block_scan_callback; + typename BlockBLevTileCountScanT::TempStorage block_scan_storage; + } blev; + + // Stage 3.3: run-length decode & block exchange for tlev + // rld_state needs to be persistent across loop iterations (RunLengthDecode calls) and, + // hence, cannot alias block_exchange_storage + struct + { + typename BlockRunLengthDecodeT::TempStorage rld_state; + typename BlockExchangeTLevT::TempStorage block_exchange_storage; + } tlev; + }; + } staged; + }; + BufferOffsetT blev_buffer_offset; + }; + + //----------------------------------------------------------------------------- + // PUBLIC TYPE MEMBERS + //----------------------------------------------------------------------------- + +public: + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //----------------------------------------------------------------------------- + // PRIVATE MEMBER FUNCTIONS + //----------------------------------------------------------------------------- + +private: + /// Shared storage reference + _TempStorage& temp_storage; + + /** + * @brief Loads this tile's buffers' sizes, without any guards (i.e., out-of-bounds checks) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void + LoadBufferSizesFullTile(BufferSizeIteratorT tile_buffer_sizes_it, BufferSizeT (&buffer_sizes)[BUFFERS_PER_THREAD]) + { + BufferLoadT(temp_storage.load_storage).Load(tile_buffer_sizes_it, buffer_sizes); + } + + /** + * @brief Loads this tile's buffers' sizes, making sure to read at most \p num_valid items. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadBufferSizesPartialTile( + BufferSizeIteratorT tile_buffer_sizes_it, BufferSizeT (&buffer_sizes)[BUFFERS_PER_THREAD], BufferOffsetT num_valid) + { + // Out-of-bounds buffer items are initialized to '0', so those buffers will simply be ignored + // later on + constexpr BufferSizeT OOB_DEFAULT_BUFFER_SIZE = 0U; + + BufferLoadT(temp_storage.load_storage).Load(tile_buffer_sizes_it, buffer_sizes, num_valid, OOB_DEFAULT_BUFFER_SIZE); + } + + /** + * @brief Computes the histogram over the number of buffers belonging to each of the three + * size-classes (TLEV, WLEV, BLEV). + */ + _CCCL_DEVICE _CCCL_FORCEINLINE VectorizedSizeClassCounterT + GetBufferSizeClassHistogram(const BufferSizeT (&buffer_sizes)[BUFFERS_PER_THREAD]) + { + VectorizedSizeClassCounterT vectorized_counters{}; + + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < BUFFERS_PER_THREAD; i++) + { + // Whether to increment ANY of the buffer size classes at all + const uint32_t increment = buffer_sizes[i] > 0 ? 1U : 0U; + // Identify the buffer's size class + uint32_t buffer_size_class = 0; + buffer_size_class += buffer_sizes[i] > WARP_LEVEL_THRESHOLD ? 1U : 0U; + buffer_size_class += buffer_sizes[i] > BLOCK_LEVEL_THRESHOLD ? 1U : 0U; + + // Increment the count of the respective size class + vectorized_counters.add(buffer_size_class, increment); + } + return vectorized_counters; + } + + /** + * @brief Scatters the buffers into the respective buffer's size-class partition. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void PartitionBuffersBySize( + const BufferSizeT (&buffer_sizes)[BUFFERS_PER_THREAD], + VectorizedSizeClassCounterT& vectorized_offsets, + BufferTuple (&buffers_by_size_class)[BUFFERS_PER_BLOCK]) + { + // If we intend to perform a stable partitioning, the thread's buffer are in a blocked + // arrangement, otherwise they are in a striped arrangement + BlockBufferOffsetT buffer_id = BUFFER_STABLE_PARTITION ? (BUFFERS_PER_THREAD * threadIdx.x) : (threadIdx.x); + constexpr BlockBufferOffsetT BUFFER_STRIDE = + BUFFER_STABLE_PARTITION ? static_cast(1) : static_cast(BLOCK_THREADS); + + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < BUFFERS_PER_THREAD; i++) + { + if (buffer_sizes[i] > 0) + { + uint32_t buffer_size_class = 0; + buffer_size_class += buffer_sizes[i] > WARP_LEVEL_THRESHOLD ? 1U : 0U; + buffer_size_class += buffer_sizes[i] > BLOCK_LEVEL_THRESHOLD ? 1U : 0U; + const uint32_t write_offset = vectorized_offsets.get(buffer_size_class); + buffers_by_size_class[write_offset] = {static_cast(buffer_sizes[i]), buffer_id}; + vectorized_offsets.add(buffer_size_class, 1U); + } + buffer_id += BUFFER_STRIDE; + } + } + + /** + * @brief Read in all the buffers that require block-level collaboration and put them to a queue + * that will get picked up in a separate, subsequent kernel. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void EnqueueBLEVBuffers( + BufferTuple* buffers_by_size_class, + InputBufferIt tile_buffer_srcs, + OutputBufferIt tile_buffer_dsts, + BufferSizeIteratorT tile_buffer_sizes, + BlockBufferOffsetT num_blev_buffers, + BufferOffsetT tile_buffer_offset, + BufferOffsetT tile_id) + { + BlockOffsetT block_offset[BLEV_BUFFERS_PER_THREAD]; + // Read in the BLEV buffer partition (i.e., the buffers that require block-level collaboration) + uint32_t blev_buffer_offset = threadIdx.x * BLEV_BUFFERS_PER_THREAD; + + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < BLEV_BUFFERS_PER_THREAD; i++) + { + if (blev_buffer_offset < num_blev_buffers) + { + BlockBufferOffsetT tile_buffer_id = buffers_by_size_class[blev_buffer_offset].buffer_id; + block_offset[i] = ::cuda::ceil_div(+tile_buffer_sizes[tile_buffer_id], BLOCK_LEVEL_TILE_SIZE); + } + else + { + // Out-of-bounds buffers are assigned a tile count of '0' + block_offset[i] = 0U; + } + blev_buffer_offset++; + } + + if (tile_id == 0) + { + BlockOffsetT block_aggregate; + BlockBLevTileCountScanT(temp_storage.staged.blev.block_scan_storage) + .ExclusiveSum(block_offset, block_offset, block_aggregate); + if (threadIdx.x == 0) + { + blev_block_scan_state.SetInclusive(0, block_aggregate); + } + } + else + { + BLevBlockScanPrefixCallbackOpT blev_tile_prefix_op( + blev_block_scan_state, temp_storage.staged.blev.block_scan_callback, ::cuda::std::plus<>{}, tile_id); + BlockBLevTileCountScanT(temp_storage.staged.blev.block_scan_storage) + .ExclusiveSum(block_offset, block_offset, blev_tile_prefix_op); + } + __syncthreads(); + + // Read in the BLEV buffer partition (i.e., the buffers that require block-level collaboration) + blev_buffer_offset = threadIdx.x * BLEV_BUFFERS_PER_THREAD; + + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < BLEV_BUFFERS_PER_THREAD; i++) + { + if (blev_buffer_offset < num_blev_buffers) + { + BlockBufferOffsetT tile_buffer_id = buffers_by_size_class[blev_buffer_offset].buffer_id; + blev_buffer_srcs[tile_buffer_offset + blev_buffer_offset] = tile_buffer_srcs[tile_buffer_id]; + blev_buffer_dsts[tile_buffer_offset + blev_buffer_offset] = tile_buffer_dsts[tile_buffer_id]; + blev_buffer_sizes[tile_buffer_offset + blev_buffer_offset] = tile_buffer_sizes[tile_buffer_id]; + blev_buffer_tile_offsets[tile_buffer_offset + blev_buffer_offset] = block_offset[i]; + blev_buffer_offset++; + } + } + } + + /** + * @brief Read in all the buffers of this tile that require warp-level collaboration and copy + * their bytes to the corresponding destination buffer + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void BatchMemcpyWLEVBuffers( + BufferTuple* buffers_by_size_class, + InputBufferIt tile_buffer_srcs, + OutputBufferIt tile_buffer_dsts, + BufferSizeIteratorT tile_buffer_sizes, + BlockBufferOffsetT num_wlev_buffers) + { + const int32_t warp_id = static_cast(threadIdx.x / warp_threads); + constexpr uint32_t warps_per_block = BLOCK_THREADS / warp_threads; + + for (BlockBufferOffsetT buffer_offset = warp_id; buffer_offset < num_wlev_buffers; buffer_offset += warps_per_block) + { + const auto buffer_id = buffers_by_size_class[buffer_offset].buffer_id; + copy_items( + tile_buffer_srcs[buffer_id], tile_buffer_dsts[buffer_id], tile_buffer_sizes[buffer_id]); + } + } + + /** + * @brief Read in all the buffers of this tile that require thread-level collaboration and copy + * their bytes to the corresponding destination buffer + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void BatchMemcpyTLEVBuffers( + BufferTuple* buffers_by_size_class, + InputBufferIt tile_buffer_srcs, + OutputBufferIt tile_buffer_dsts, + BlockBufferOffsetT num_tlev_buffers) + { + // Read in the buffers' ids that require thread-level collaboration (where buffer id is the + // buffer within this tile) + BlockBufferOffsetT tlev_buffer_ids[TLEV_BUFFERS_PER_THREAD]; + TLevBufferSizeT tlev_buffer_sizes[TLEV_BUFFERS_PER_THREAD]; + // Currently we do not go over the TLEV buffers in multiple iterations, so we need to make sure + // we are able to be covered for the case that all our buffers are TLEV buffers + static_assert(TLEV_BUFFERS_PER_THREAD >= BUFFERS_PER_THREAD, + "Unsupported confiugraiton: The number of 'thread-level buffers' must be at " + "least as large as the number of overall buffers being processed by each " + "thread."); + + // Read in the TLEV buffer partition (i.e., the buffers that require thread-level collaboration) + uint32_t tlev_buffer_offset = threadIdx.x * TLEV_BUFFERS_PER_THREAD; + + // Pre-populate the buffer sizes to 0 (i.e. zero-padding towards the end) to ensure + // out-of-bounds TLEV buffers will not be considered + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < TLEV_BUFFERS_PER_THREAD; i++) + { + tlev_buffer_sizes[i] = 0; + } + + // Assign TLEV buffers in a blocked arrangement (each thread is assigned consecutive TLEV + // buffers) + _CCCL_PRAGMA_UNROLL_FULL() + for (uint32_t i = 0; i < TLEV_BUFFERS_PER_THREAD; i++) + { + if (tlev_buffer_offset < num_tlev_buffers) + { + tlev_buffer_ids[i] = buffers_by_size_class[tlev_buffer_offset].buffer_id; + tlev_buffer_sizes[i] = buffers_by_size_class[tlev_buffer_offset].size; + } + tlev_buffer_offset++; + } + + // Evenly distribute all the bytes that have to be copied from all the buffers that require + // thread-level collaboration using BlockRunLengthDecode + uint32_t num_total_tlev_bytes = 0U; + BlockRunLengthDecodeT block_run_length_decode( + temp_storage.staged.tlev.rld_state, tlev_buffer_ids, tlev_buffer_sizes, num_total_tlev_bytes); + + // Run-length decode the buffers' sizes into a window buffer of limited size. This is repeated + // until we were able to cover all the bytes of TLEV buffers + uint32_t decoded_window_offset = 0U; + while (decoded_window_offset < num_total_tlev_bytes) + { + BlockBufferOffsetT buffer_id[TLEV_BYTES_PER_THREAD]; + TLevBufferSizeT buffer_byte_offset[TLEV_BYTES_PER_THREAD]; + + // Now we have a balanced assignment: buffer_id[i] will hold the tile's buffer id and + // buffer_byte_offset[i] that buffer's byte that this thread supposed to copy + block_run_length_decode.RunLengthDecode(buffer_id, buffer_byte_offset, decoded_window_offset); + + // Zip from SoA to AoS + ZippedTLevByteAssignment zipped_byte_assignment[TLEV_BYTES_PER_THREAD]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int32_t i = 0; i < TLEV_BYTES_PER_THREAD; i++) + { + zipped_byte_assignment[i] = {buffer_id[i], buffer_byte_offset[i]}; + } + + // Exchange from blocked to striped arrangement for coalesced memory reads and writes + BlockExchangeTLevT(temp_storage.staged.tlev.block_exchange_storage) + .BlockedToStriped(zipped_byte_assignment, zipped_byte_assignment); + + // Read in the bytes that this thread is assigned to + constexpr uint32_t WINDOW_SIZE = (TLEV_BYTES_PER_THREAD * BLOCK_THREADS); + const bool is_full_window = decoded_window_offset + WINDOW_SIZE < num_total_tlev_bytes; + if (is_full_window) + { + uint32_t absolute_tlev_byte_offset = decoded_window_offset + threadIdx.x; + AliasT src_byte[TLEV_BYTES_PER_THREAD]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int32_t i = 0; i < TLEV_BYTES_PER_THREAD; i++) + { + src_byte[i] = read_item( + tile_buffer_srcs[zipped_byte_assignment[i].tile_buffer_id], zipped_byte_assignment[i].buffer_byte_offset); + absolute_tlev_byte_offset += BLOCK_THREADS; + } + + _CCCL_PRAGMA_UNROLL_FULL() + for (int32_t i = 0; i < TLEV_BYTES_PER_THREAD; i++) + { + write_item( + tile_buffer_dsts[zipped_byte_assignment[i].tile_buffer_id], + zipped_byte_assignment[i].buffer_byte_offset, + src_byte[i]); + } + } + else + { + uint32_t absolute_tlev_byte_offset = decoded_window_offset + threadIdx.x; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int32_t i = 0; i < TLEV_BYTES_PER_THREAD; i++) + { + if (absolute_tlev_byte_offset < num_total_tlev_bytes) + { + const AliasT src_byte = read_item( + tile_buffer_srcs[zipped_byte_assignment[i].tile_buffer_id], zipped_byte_assignment[i].buffer_byte_offset); + write_item( + tile_buffer_dsts[zipped_byte_assignment[i].tile_buffer_id], + zipped_byte_assignment[i].buffer_byte_offset, + src_byte); + } + absolute_tlev_byte_offset += BLOCK_THREADS; + } + } + + decoded_window_offset += WINDOW_SIZE; + + // Ensure all threads finished collaborative BlockExchange so temporary storage can be reused + // with next iteration + __syncthreads(); + } + } + + //----------------------------------------------------------------------------- + // PUBLIC MEMBER FUNCTIONS + //----------------------------------------------------------------------------- + +public: + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTile(BufferOffsetT tile_id) + { + // Offset into this tile's buffers + BufferOffsetT buffer_offset = tile_id * BUFFERS_PER_BLOCK; + + // Indicates whether all of this tiles items are within bounds + bool is_full_tile = buffer_offset + BUFFERS_PER_BLOCK < num_buffers; + + // Load the buffer sizes of this tile's buffers + BufferSizeIteratorT tile_buffer_sizes_it = buffer_sizes_it + buffer_offset; + BufferSizeT buffer_sizes[BUFFERS_PER_THREAD]; + if (is_full_tile) + { + LoadBufferSizesFullTile(tile_buffer_sizes_it, buffer_sizes); + } + else + { + LoadBufferSizesPartialTile(tile_buffer_sizes_it, buffer_sizes, num_buffers - buffer_offset); + } + + // Ensure we can repurpose the BlockLoad's temporary storage + __syncthreads(); + + // Count how many buffers fall into each size-class + VectorizedSizeClassCounterT size_class_histogram = GetBufferSizeClassHistogram(buffer_sizes); + + // Compute the prefix sum over the histogram + VectorizedSizeClassCounterT size_class_agg = {}; + BlockSizeClassScanT(temp_storage.size_scan_storage) + .ExclusiveSum(size_class_histogram, size_class_histogram, size_class_agg); + + // Ensure we can repurpose the scan's temporary storage for scattering the buffer ids + __syncthreads(); + + // Factor in the per-size-class counts / offsets + // That is, WLEV buffer offset has to be offset by the TLEV buffer count and BLEV buffer offset + // has to be offset by the TLEV+WLEV buffer count + uint32_t buffer_count = 0U; + for (uint32_t i = 0; i < NUM_SIZE_CLASSES; i++) + { + size_class_histogram.add(i, buffer_count); + buffer_count += size_class_agg.get(i); + } + + // Signal the number of BLEV buffers we're planning to write out + BufferOffsetT buffer_exclusive_prefix = 0; + if (tile_id == 0) + { + if (threadIdx.x == 0) + { + blev_buffer_scan_state.SetInclusive(tile_id, size_class_agg.get(BLEV_SIZE_CLASS)); + } + buffer_exclusive_prefix = 0; + } + else + { + BLevBuffScanPrefixCallbackOpT blev_buffer_prefix_op( + blev_buffer_scan_state, temp_storage.buffer_scan_callback, ::cuda::std::plus<>{}, tile_id); + + // Signal our partial prefix and wait for the inclusive prefix of previous tiles + if (threadIdx.x < warp_threads) + { + buffer_exclusive_prefix = blev_buffer_prefix_op(size_class_agg.get(BLEV_SIZE_CLASS)); + } + } + if (threadIdx.x == 0) + { + temp_storage.blev_buffer_offset = buffer_exclusive_prefix; + } + + // Ensure the prefix callback has finished using its temporary storage and that it can be reused + // in the next stage + __syncthreads(); + + // Scatter the buffers into one of the three partitions (TLEV, WLEV, BLEV) depending on their + // size + PartitionBuffersBySize(buffer_sizes, size_class_histogram, temp_storage.staged.buffers_by_size_class); + + // Ensure all buffers have been partitioned by their size class AND + // ensure that blev_buffer_offset has been written to shared memory + __syncthreads(); + + // TODO: think about prefetching tile_buffer_{srcs,dsts} into shmem + InputBufferIt tile_buffer_srcs = input_buffer_it + buffer_offset; + OutputBufferIt tile_buffer_dsts = output_buffer_it + buffer_offset; + BufferSizeIteratorT tile_buffer_sizes = buffer_sizes_it + buffer_offset; + + // Copy block-level buffers + EnqueueBLEVBuffers( + &temp_storage.staged + .buffers_by_size_class[size_class_agg.get(TLEV_SIZE_CLASS) + size_class_agg.get(WLEV_SIZE_CLASS)], + tile_buffer_srcs, + tile_buffer_dsts, + tile_buffer_sizes, + size_class_agg.get(BLEV_SIZE_CLASS), + temp_storage.blev_buffer_offset, + tile_id); + + // Ensure we can repurpose the temporary storage required by EnqueueBLEVBuffers + __syncthreads(); + + // Copy warp-level buffers + BatchMemcpyWLEVBuffers( + &temp_storage.staged.buffers_by_size_class[size_class_agg.get(TLEV_SIZE_CLASS)], + tile_buffer_srcs, + tile_buffer_dsts, + tile_buffer_sizes, + size_class_agg.get(WLEV_SIZE_CLASS)); + + // Perform batch memcpy for all the buffers that require thread-level collaboration + uint32_t num_tlev_buffers = size_class_agg.get(TLEV_SIZE_CLASS); + BatchMemcpyTLEVBuffers( + temp_storage.staged.buffers_by_size_class, tile_buffer_srcs, tile_buffer_dsts, num_tlev_buffers); + } + + //----------------------------------------------------------------------------- + // CONSTRUCTOR + //----------------------------------------------------------------------------- + _CCCL_DEVICE _CCCL_FORCEINLINE AgentBatchMemcpy( + TempStorage& temp_storage, + InputBufferIt input_buffer_it, + OutputBufferIt output_buffer_it, + BufferSizeIteratorT buffer_sizes_it, + BufferOffsetT num_buffers, + BlevBufferSrcsOutItT blev_buffer_srcs, + BlevBufferDstsOutItT blev_buffer_dsts, + BlevBufferSizesOutItT blev_buffer_sizes, + BlevBufferTileOffsetsOutItT blev_buffer_tile_offsets, + BLevBufferOffsetTileState blev_buffer_scan_state, + BLevBlockOffsetTileState blev_block_scan_state) + : temp_storage(temp_storage.Alias()) + , input_buffer_it(input_buffer_it) + , output_buffer_it(output_buffer_it) + , buffer_sizes_it(buffer_sizes_it) + , num_buffers(num_buffers) + , blev_buffer_srcs(blev_buffer_srcs) + , blev_buffer_dsts(blev_buffer_dsts) + , blev_buffer_sizes(blev_buffer_sizes) + , blev_buffer_tile_offsets(blev_buffer_tile_offsets) + , blev_buffer_scan_state(blev_buffer_scan_state) + , blev_block_scan_state(blev_block_scan_state) + {} + +private: + // Iterator providing the pointers to the source memory buffers + InputBufferIt input_buffer_it; + // Iterator providing the pointers to the destination memory buffers + OutputBufferIt output_buffer_it; + // Iterator providing the number of bytes to be copied for each pair of buffers + BufferSizeIteratorT buffer_sizes_it; + // The total number of buffer pairs + BufferOffsetT num_buffers; + // Output iterator to which the source pointers of the BLEV buffers are written + BlevBufferSrcsOutItT blev_buffer_srcs; + // Output iterator to which the destination pointers of the BLEV buffers are written + BlevBufferDstsOutItT blev_buffer_dsts; + // Output iterator to which the number of bytes to be copied of the BLEV buffers are written + BlevBufferSizesOutItT blev_buffer_sizes; + // Output iterator to which the mapping of tiles to BLEV buffers is written + BlevBufferTileOffsetsOutItT blev_buffer_tile_offsets; + // The single-pass prefix scan's tile state used for tracking the prefix sum over the number of + // BLEV buffers + BLevBufferOffsetTileState blev_buffer_scan_state; + // The single-pass prefix scan's tile state used for tracking the prefix sum over tiles of BLEV + // buffers + BLevBlockOffsetTileState blev_block_scan_state; +}; +} // namespace detail::batch_memcpy + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_batched_topk.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_batched_topk.cuh new file mode 100644 index 00000000..d2a99cfc --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_batched_topk.cuh @@ -0,0 +1,375 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::batched_topk +{ +// Atomic counters used by the small-segment kernel to (a) enqueue large segments into the large-segment work queue +// and (b) elect the last block to run the epilogue scan over the queued tile counts. `alignas(128)` isolates each +// counter on its own cache line for performance. +template +struct batched_topk_counters +{ + // Force unsigned integer type for segment count. + using segment_count_t = detail::choose_offset_t; + // Number of segments enqueued in the large-segment work queue. Atomically incremented (by 1) by the first thread + // of each block that decides its segment is large. + alignas(128) segment_count_t large_segments_count; + + // Block retirement counter. Each block atomically increments by 1 when it has finished processing its segment, and + // the block that observes `gridDim.x - 1` runs the epilogue on the queued large segments tile counts. + // Assumption: Future support for more than 2^31 - 1 segments will use multiple launches of a slightly modified + // small-segment kernel instead of additional grid dimensions. Therefore each grid will handle a maximum of 2^31 - 1 + // segments per launch. The counter would not even have to be reset to 0 after each launch if we cleverly make use of + // its modulo arithmetic. + alignas(128) unsigned retirement_count; +}; + +template +struct agent_batched_topk_worker_per_segment +{ + // ------------------------------------------------------------------------- + // Types and Constants + // ------------------------------------------------------------------------- + // Derive inner types from Iterator of Iterators + using key_it_t = it_value_t; + using value_it_t = it_value_t; + + using key_t = it_value_t; + using value_t = it_value_t; + + using segment_size_val_t = typename ::cuda::args::__traits::element_type; + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + using counters_t = batched_topk_counters; + + static constexpr auto policy = PolicyGetter{}(); + static constexpr worker_policy active_policy = policy.worker_per_segment_policy; + + // For block-topk (and keys/values load/store): + static constexpr int threads_per_block = active_policy.threads_per_block; + static constexpr int items_per_thread = active_policy.items_per_thread; + static constexpr int tile_size = threads_per_block * items_per_thread; + + // For block-scan (and offsets load/store): + static constexpr int epilogue_items_per_thread = active_policy.epilogue.items_per_thread; + static constexpr int epilogue_tile_size = threads_per_block * epilogue_items_per_thread; + + // Number used for preprocessing segment-size data, not for tuning => should not affect performance of this agent. + static constexpr multi_worker_policy multi_worker_per_segment_policy = policy.multi_worker_per_segment_policy; + static constexpr int multi_worker_per_segment_tile_size = + multi_worker_per_segment_policy.threads_per_block * multi_worker_per_segment_policy.items_per_thread; + + // Check if there could be large segments present + static constexpr bool only_small_segments = ::cuda::args::__traits::highest <= tile_size; + + // Check if we are dealing with keys-only or key-value pairs + static constexpr bool is_keys_only = ::cuda::std::is_same_v; + + // ------------------------------------------------------------------------- + // Primitive Types + // ------------------------------------------------------------------------- + using block_load_keys_t = BlockLoad; + using block_load_vals_t = BlockLoad; + + using block_topk_t = block_topk; + + // TODO (elstehle): Specialize for the case that we statically know k and we can skip passing num_valid_items to + // Store() + using block_store_keys_t = BlockStore; + using block_store_vals_t = BlockStore; + + using block_load_epilogue_t = + BlockLoad; + using block_scan_epilogue_t = BlockScan; + using block_store_epilogue_t = + BlockStore; + + // ------------------------------------------------------------------------- + // Shared Memory Storage + // ------------------------------------------------------------------------- + struct TempStorage_ + { + union + { + typename block_load_keys_t::TempStorage load_keys; + typename block_load_vals_t::TempStorage load_vals; + typename block_topk_t::TempStorage topk; + typename block_store_keys_t::TempStorage store_keys; + typename block_store_vals_t::TempStorage store_vals; + typename block_load_epilogue_t::TempStorage load_epilogue; + typename block_scan_epilogue_t::TempStorage scan_epilogue; + typename block_store_epilogue_t::TempStorage store_epilogue; + }; + }; + + using TempStorage = Uninitialized; + + // ------------------------------------------------------------------------- + // Members + // ------------------------------------------------------------------------- + TempStorage_& temp_storage; + KeyInputItItT d_key_segments_it; + KeyOutputItItT d_key_segments_out_it; + ValueInputItItT d_value_segments_it; + ValueOutputItItT d_value_segments_out_it; + SegmentSizeParameterT segment_sizes; + KParameterT k_param; + SelectDirectionParameterT select_directions; + NumSegmentsParameterT num_segments; + counters_t* d_counters; + num_segments_val_t* d_large_segments_ids; + LargeSegmentTileOffsetT* d_large_segments_tile_offsets; + // ------------------------------------------------------------------------- + // Constructor + // ------------------------------------------------------------------------- + _CCCL_DEVICE_API _CCCL_FORCEINLINE agent_batched_topk_worker_per_segment( + TempStorage& temp_storage, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k_param, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + counters_t* d_counters, + num_segments_val_t* d_large_segments_ids, + LargeSegmentTileOffsetT* d_large_segments_tile_offsets) + : temp_storage(temp_storage.Alias()) + , d_key_segments_it(d_key_segments_it) + , d_key_segments_out_it(d_key_segments_out_it) + , d_value_segments_it(d_value_segments_it) + , d_value_segments_out_it(d_value_segments_out_it) + , segment_sizes(segment_sizes) + , k_param(k_param) + , select_directions(select_directions) + , num_segments(num_segments) + , d_counters(d_counters) + , d_large_segments_ids(d_large_segments_ids) + , d_large_segments_tile_offsets(d_large_segments_tile_offsets) + {} + + _CCCL_DEVICE_API _CCCL_FORCEINLINE void Process() + { + // Identify Segment + const int segment_id = static_cast(blockIdx.x); + + // Boundary check + // TODO (elstehle): consider skipping boundary check if we can safely assume the right grid dimensions + if (segment_id >= params::get_param(num_segments, 0)) + { + return; + } + + constexpr bool is_full_tile = ::cuda::args::__traits::is_constant + && ::cuda::args::__traits::lowest == tile_size; + + // Resolve Segment Parameters + const auto segment_size = params::get_param(segment_sizes, segment_id); + if (!only_small_segments && segment_size > tile_size) + { + // Enqueue large segment + if (threadIdx.x == 0u) + { + // Add to large segment queue + const auto large_segment_queue_idx = atomicAdd(&d_counters->large_segments_count, 1ull); + d_large_segments_ids[large_segment_queue_idx] = static_cast(segment_id); + d_large_segments_tile_offsets[large_segment_queue_idx] = + static_cast(::cuda::ceil_div(segment_size, multi_worker_per_segment_tile_size)); + } + } + else + { + // Process small segment + const auto k = (::cuda::std::min) (params::get_param(k_param, segment_id), + static_cast(segment_size)); + const auto direction = select_directions.get_param(segment_id); + + // Determine padding key based on direction + const key_t padding_key = + (direction == detail::topk::select::max) + ? ::cuda::std::numeric_limits::lowest() + : (::cuda::std::numeric_limits::max)(); + + // Dereference iterator-of-iterators to get the segment specific iterator + auto block_keys_in = d_key_segments_it[segment_id]; + + // Load Keys + key_t thread_keys[items_per_thread]; + if constexpr (is_full_tile) + { + // No padding needed + block_load_keys_t(temp_storage.load_keys).Load(block_keys_in, thread_keys); + } + else + { + // Potentially partial final load with padding + // TODO (elstehle): explore whether a runtime check for segment_size == tile_size improves performance + block_load_keys_t(temp_storage.load_keys).Load(block_keys_in, thread_keys, segment_size); + } + + // Load Values (if applicable) + [[maybe_unused]] value_t thread_values[items_per_thread]; + + if constexpr (!is_keys_only) + { + __syncthreads(); + auto block_vals_in = d_value_segments_it[segment_id]; + + if constexpr (is_full_tile) + { + // No padding needed + block_load_vals_t(temp_storage.load_vals).Load(block_vals_in, thread_values); + } + else + { + // Potentially partial final load with padding + // TODO (elstehle): explore whether a runtime check for segment_size == tile_size improves performance + block_load_vals_t(temp_storage.load_vals).Load(block_vals_in, thread_values, segment_size); + } + } + + __syncthreads(); + + // Perform Block Top-K + if constexpr (is_keys_only) + { + const bool is_successful_dispatch = cub::detail::params::dispatch_discrete( + select_directions, segment_id, [this, &thread_keys, k, segment_size](auto direction_tag) { + if constexpr (decltype(direction_tag)::value == detail::topk::select::max) + { + block_topk_t(temp_storage.topk).template max_keys(thread_keys, k, segment_size); + } + else + { + block_topk_t(temp_storage.topk).template min_keys(thread_keys, k, segment_size); + } + }); + _CCCL_ASSERT(is_successful_dispatch, "Error: Unsupported select direction"); + } + else + { + // Pass both keys and values + const bool is_successful_dispatch = cub::detail::params::dispatch_discrete( + select_directions, segment_id, [this, &thread_keys, &thread_values, k, segment_size](auto direction_tag) { + if constexpr (decltype(direction_tag)::value == detail::topk::select::max) + { + block_topk_t(temp_storage.topk) + .template max_pairs(thread_keys, thread_values, k, segment_size); + } + else + { + block_topk_t(temp_storage.topk) + .template min_pairs(thread_keys, thread_values, k, segment_size); + } + }); + _CCCL_ASSERT(is_successful_dispatch, "Error: Unsupported select direction"); + } + + __syncthreads(); + + auto block_keys_out = d_key_segments_out_it[segment_id]; + + block_store_keys_t(temp_storage.store_keys) + .Store(block_keys_out, + thread_keys, + k // Only store K items + ); + + if constexpr (!is_keys_only) + { + __syncthreads(); + auto block_vals_out = d_value_segments_out_it[segment_id]; + + block_store_vals_t(temp_storage.store_vals).Store(block_vals_out, thread_values, k); + } + } + + // Epilogue: Scan queued large segment sizes (in tiles not elements) for load balancing search in the large segment + // agent + if constexpr (!only_small_segments) + { + // Determine last block trying to retire. + bool is_last_block = false; + if (threadIdx.x == 0u) + { + __threadfence(); + const auto retirement_count = atomicAdd(&d_counters->retirement_count, 1u); + is_last_block = retirement_count == (gridDim.x - 1u); + } + // This sync also makes sure that the shared memory can be reused. + is_last_block = static_cast(__syncthreads_or(static_cast(is_last_block))); + if (!is_last_block) + { + return; + } + const auto num_large_segments = d_counters->large_segments_count; + // For tracking the running total across tiles (loop iterations). + // Caution: The functor is only invoked by the first warp in the block, and the value returned by lane 0 in that + // warp is used as the initial value. + const auto prefix_callback_op = + [running_total = segment_size_val_t{0}](segment_size_val_t block_aggregate) mutable { + auto old_running_total = running_total; + running_total += block_aggregate; + return old_running_total; + }; + _CCCL_PRAGMA_NOUNROLL() + for (int large_segment_offset = 0; large_segment_offset < num_large_segments; + large_segment_offset += epilogue_tile_size) + { + segment_size_val_t segment_tile_offsets[epilogue_items_per_thread]; + block_load_epilogue_t(temp_storage.load_epilogue) + .Load(d_large_segments_tile_offsets + large_segment_offset, + segment_tile_offsets, + num_large_segments - large_segment_offset, + 0); + __syncthreads(); + block_scan_epilogue_t(temp_storage.scan_epilogue) + .ExclusiveSum(segment_tile_offsets, segment_tile_offsets, prefix_callback_op); + __syncthreads(); + block_store_epilogue_t(temp_storage.store_epilogue) + .Store(d_large_segments_tile_offsets + large_segment_offset, + segment_tile_offsets, + num_large_segments - large_segment_offset); + __syncthreads(); + } + } + } +}; +} // namespace detail::batched_topk +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_find.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_find.cuh new file mode 100644 index 00000000..33265175 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_find.cuh @@ -0,0 +1,201 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once +#include + +#include +#include +#include +#include + +#include +#include + +#include +#if !_CCCL_HAS_NV_ATOMIC_BUILTINS() +# include +#endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS() +#include + +CUB_NAMESPACE_BEGIN +namespace detail::find +{ +template +struct agent_t +{ + // The input value type + using InputT = typename ::cuda::std::iterator_traits::value_type; + + // Vector type of InputT for data movement + using VectorT = typename CubVector::Type; + + static constexpr int tile_size = ThreadsPerBlock * ItemsPerThread; + + // Can vectorize according to the policy if the input iterator is a native pointer to a primitive type + static constexpr bool attempt_vectorization = + (VecSize > 1) && (ItemsPerThread % VecSize == 0) && (::cuda::std::contiguous_iterator) + && THRUST_NS_QUALIFIER::is_trivially_relocatable_v; + + static constexpr CacheLoadModifier load_modifier = LoadModifier; + + // Shared memory type required by this thread block + struct _TempStorage + { + OffsetT global_result; + OffsetT block_result; + }; + + // Alias wrapper allowing storage to be unioned + using TempStorage = Uninitialized<_TempStorage>; + + _TempStorage& temp_storage; + InputIteratorT d_in; + PredicateT predicate; + OffsetT* found_pos_ptr; + OffsetT num_items; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE bool is_aligned_and_full_tile(OffsetT tile_offset) + { + if constexpr (CanVectorize) + { + static_assert(::cuda::std::is_pointer_v); + + // Retrieve the value type from the iterator to determine the vector type + using InputT = typename ::cuda::std::iterator_traits::value_type; + using VectorT = typename CubVector::Type; + + const bool full_tile = (tile_offset + tile_size) <= num_items; + + // Check alignment at the actual load position (d_in + tile_offset) + return full_tile && ::cuda::is_aligned(d_in + tile_offset, sizeof(VectorT)); + } + else + { + return false; + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE bool + ConsumeTile(OffsetT tile_offset, ::cuda::std::integral_constant /*CAN_VECTORIZE*/) + { + using InputT = typename ::cuda::std::iterator_traits::value_type; + using VectorT = typename CubVector::Type; + + // vectorized loads begin + auto load_ptr = reinterpret_cast(d_in + tile_offset + (threadIdx.x * VecSize)); + CacheModifiedInputIterator d_vec_in(load_ptr); + + alignas(InputT) unsigned char input_bytes[ItemsPerThread * sizeof(InputT)]; + auto* vec_items = reinterpret_cast(input_bytes); + + constexpr int number_of_vectors = ItemsPerThread / VecSize; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < number_of_vectors; ++i) + { + vec_items[i] = d_vec_in[ThreadsPerBlock * i]; + } + + for (int i = 0; i < ItemsPerThread; ++i) + { + OffsetT nth_vector_of_thread = i / VecSize; + OffsetT element_in_vector = i % VecSize; + OffsetT vector_of_tile = nth_vector_of_thread * ThreadsPerBlock + threadIdx.x; + + OffsetT index = tile_offset + vector_of_tile * VecSize + element_in_vector; + + auto* input_items = reinterpret_cast(input_bytes); + if (index < num_items && predicate(input_items[i])) + { + atomicMin(&temp_storage.block_result, index); + // every thread goes over multiple elements per thread for every tile. If a thread finds a local minimum it + // doesn't need to proceed further (inner early exit). + return true; + } + } + + return false; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE bool + ConsumeTile(OffsetT tile_offset, ::cuda::std::integral_constant /*CAN_VECTORIZE*/) + { + for (int i = 0; i < ItemsPerThread; ++i) + { + const auto index = tile_offset + threadIdx.x + i * blockDim.x; + if (index < num_items) + { + // using raw_reference_cast and passing directly to predicate should avoid creating a copy, and thus prevent + // bugs like: http://github.com/NVIDIA/cccl/issues/3591 + if (predicate(THRUST_NS_QUALIFIER::raw_reference_cast(d_in[index]))) + { + atomicMin(&temp_storage.block_result, index); + return true; + } + } + } + + return false; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void Process() + { + if (threadIdx.x == 0) + { + temp_storage.block_result = num_items; + } + __syncthreads(); + + // use a grid strided loop + OffsetT grid_stride = static_cast(tile_size) * static_cast(gridDim.x); + for (OffsetT tile_offset = static_cast(blockIdx.x) * static_cast(tile_size); + tile_offset < num_items; + tile_offset += grid_stride) + { + // Only one thread reads atomically and propagates it to other threads of the block through shared memory + if (threadIdx.x == 0) + { +#if _CCCL_HAS_NV_ATOMIC_BUILTINS() + // __nv_atomic_load is a compiler build-in and compiles a lot faster + __nv_atomic_load(found_pos_ptr, &temp_storage.global_result, __NV_ATOMIC_RELAXED, __NV_THREAD_SCOPE_DEVICE); +#else // ^^^ _CCCL_HAS_NV_ATOMIC_BUILTINS() ^^^ / vvv !_CCCL_HAS_NV_ATOMIC_BUILTINS() vvv + temp_storage.global_result = ::cuda::atomic_ref{*found_pos_ptr}.load( + ::cuda::std::memory_order_relaxed); +#endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS() + } + __syncthreads(); + + // early exit + if (temp_storage.global_result < tile_offset) + { + return; + } + + const bool found_thread = + is_aligned_and_full_tile(tile_offset) + ? ConsumeTile(tile_offset, ::cuda::std::bool_constant{}) + : ConsumeTile(tile_offset, ::cuda::std::false_type{}); + + const bool found_block = __syncthreads_or(found_thread); + if (found_block) + { + // our block found it, update global position and exit + if (threadIdx.x == 0 && temp_storage.block_result < num_items) + { + atomicMin(found_pos_ptr, temp_storage.block_result); + } + return; + } + } + } +}; +} // namespace detail::find + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_find_bound_sorted_values.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_find_bound_sorted_values.cuh new file mode 100644 index 00000000..78f9469f --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_find_bound_sorted_values.cuh @@ -0,0 +1,222 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::find_bound_sorted_values +{ +// lower_bound vs upper_bound: partition comparator and per-step advance differ. +struct lower_bound_mode +{ + // Wrap user comp so the merge path partitions identically to std::lower_bound. + template + struct partition_comp_t + { + CompareOp comp; + + template + _CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE bool operator()(A&& a, B&& b) const + { + return !comp(::cuda::std::forward(b), ::cuda::std::forward(a)); + } + }; + + template + _CCCL_HOST_DEVICE_API static partition_comp_t make_partition_comp(CompareOp compare_op) + { + return partition_comp_t{compare_op}; + } + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE static bool + should_advance(const HaystackT& haystack_value, const NeedlesT& needle_value, CompareOp compare_op) + { + return compare_op(haystack_value, needle_value); + } +}; + +struct upper_bound_mode +{ + template + using partition_comp_t = CompareOp; + + template + _CCCL_HOST_DEVICE_API static CompareOp make_partition_comp(CompareOp compare_op) + { + return compare_op; + } + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE static bool + should_advance(const HaystackT& haystack_value, const NeedlesT& needle_value, CompareOp compare_op) + { + return !compare_op(needle_value, haystack_value); + } +}; + +template +struct agent_t +{ + static constexpr int tile_size = ThreadsPerBlock * ItemsPerThread; + + using haystack_type = it_value_t; + using needles_type = it_value_t; + + // Separate buffers because haystack and needles may have different value types. + struct _TempStorage + { + haystack_type haystack[tile_size]; + needles_type needles[tile_size]; + }; + + using TempStorage = Uninitialized<_TempStorage>; + + _TempStorage& storage; + HaystackIt d_range; + NeedlesIt d_values; + OutputIt d_output; + Offset range_count; + Offset values_count; + Offset* range_beg_offsets; + CompareOp compare_op; + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void consume_tile(int tile_idx, Offset diag0, int total_in_tile) + { + const Offset range_beg = range_beg_offsets[tile_idx]; + const Offset range_end = range_beg_offsets[tile_idx + 1]; + _CCCL_ASSERT(range_end >= range_beg, ""); + _CCCL_ASSERT(diag0 >= range_beg, ""); + const Offset values_beg = diag0 - range_beg; + + const int haystack_count = static_cast(range_end - range_beg); + const int needles_count = total_in_tile - haystack_count; + + { + const auto d_range_cm = cub::detail::try_make_cache_modified_iterator(d_range + range_beg); + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ItemsPerThread; ++item) + { + const int idx = ThreadsPerBlock * item + threadIdx.x; + if (idx < haystack_count) + { + storage.haystack[idx] = d_range_cm[idx]; + } + } + } + + { + auto d_values_cm = cub::detail::try_make_cache_modified_iterator(d_values + values_beg); + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ItemsPerThread; ++item) + { + const int idx = ThreadsPerBlock * item + threadIdx.x; + if (idx < needles_count) + { + storage.needles[idx] = d_values_cm[idx]; + } + } + } + + __syncthreads(); + +#ifdef CCCL_ENABLE_DEVICE_ASSERTIONS + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ItemsPerThread; ++item) + { + const int idx = ThreadsPerBlock * item + threadIdx.x; + if (idx < needles_count && (values_beg + idx) > 0) + { + const needles_type prev = (idx == 0) ? d_values[values_beg - 1] : storage.needles[idx - 1]; + _CCCL_ASSERT(!compare_op(storage.needles[idx], prev), "d_values must be sorted consistently with comp"); + } + } +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS + + const auto partition_comp = Mode::make_partition_comp(compare_op); + + int d0_thread = ItemsPerThread * static_cast(threadIdx.x); + if constexpr (!IsFullTile) + { + d0_thread = ::cuda::std::min(d0_thread, total_in_tile); + } + + const int i0 = + cub::MergePath(storage.haystack, storage.needles, haystack_count, needles_count, d0_thread, partition_comp); + const int j0 = d0_thread - i0; + + int i = i0; + int j = j0; + int haystack_remaining = haystack_count - i0; + int needles_remaining = needles_count - j0; + + const int steps = IsFullTile ? ItemsPerThread : ::cuda::std::min(total_in_tile - d0_thread, ItemsPerThread); + _CCCL_PRAGMA_UNROLL(ItemsPerThread) + for (int step = 0; step < steps; ++step) + { + const bool advance_haystack = + (needles_remaining == 0) + || (haystack_remaining > 0 && Mode::should_advance(storage.haystack[i], storage.needles[j], compare_op)); + if (advance_haystack) + { + ++i; + --haystack_remaining; + } + else + { + using output_value_t = cub::detail::non_void_value_t; + d_output[values_beg + j] = static_cast(range_beg + i); + ++j; + --needles_remaining; + } + } + } + + _CCCL_DEVICE_API _CCCL_FORCEINLINE void operator()() + { + const int tile_idx = static_cast(blockIdx.x); + const Offset diag0 = static_cast(tile_size) * tile_idx; + const Offset diag1 = ::cuda::std::min(diag0 + static_cast(tile_size), range_count + values_count); + const int total_in_tile = static_cast(diag1 - diag0); + + if (total_in_tile == tile_size) + { + consume_tile(tile_idx, diag0, tile_size); + } + else + { + consume_tile(tile_idx, diag0, total_in_tile); + } + } +}; +} // namespace detail::find_bound_sorted_values + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_for.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_for.cuh new file mode 100644 index 00000000..6f159234 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_for.cuh @@ -0,0 +1,56 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::for_each +{ +template +struct policy_t +{ + static constexpr int threads_per_block = ThreadsPerBlock; + static constexpr int items_per_thread = ItemsPerThread; +}; + +template +struct agent_block_striped_t +{ + static constexpr int items_per_thread = PolicyT::items_per_thread; + + OffsetT tile_base; + OpT op; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void consume_tile(int items_in_tile, int threads_per_block) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < items_per_thread; item++) + { + const auto idx = + static_cast(threads_per_block * item + threadIdx.x); // NOLINT(bugprone-misplaced-widening-cast) + + if (IsFullTile || idx < items_in_tile) + { + (void) op(tile_base + idx); + } + } + } +}; +} // namespace detail::for_each + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_histogram.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_histogram.cuh new file mode 100644 index 00000000..daed613b --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_histogram.cuh @@ -0,0 +1,725 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! \file +//! cub::AgentHistogram implements a stateful abstraction of CUDA thread blocks for participating in device-wide +//! histogram. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +enum BlockHistogramMemoryPreference +{ + GMEM, + SMEM, + BLEND +}; + +#if _CCCL_HOSTED() +namespace detail +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(BlockHistogramMemoryPreference mempref) noexcept +{ + switch (mempref) + { + case GMEM: + return "GMEM"; + case SMEM: + return "SMEM"; + case BLEND: + return "BLEND"; + } + return ""; +} +} // namespace detail + +inline ::std::ostream& operator<<(::std::ostream& os, BlockHistogramMemoryPreference mempref) +{ + return os << CUB_NS_QUALIFIER::detail::to_string(mempref); +} +#endif // _CCCL_HOSTED() + +CUB_NAMESPACE_END + +#if __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) +template <::cuda::std::same_as CharT> +struct std::formatter : formatter +{ + template + auto format(const CUB_NS_QUALIFIER::BlockHistogramMemoryPreference& mempref, FmtCtx& ctx) const + { + return formatter::format(CUB_NS_QUALIFIER::detail::to_string(mempref), ctx); + } +}; +#endif // __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +//! Parameterizable tuning policy type for AgentHistogram +template +struct agent_histogram_policy +{ + /// Threads per thread block + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + /// Pixels per thread (per tile of input) + static constexpr int PIXELS_PER_THREAD = PixelsPerThread; + + /// Whether to perform localized RLE to compress samples before histogramming + static constexpr bool IS_RLE_COMPRESS = RleCompress; + + /// Whether to prefer privatized shared-memory bins (versus privatized global-memory bins) + static constexpr BlockHistogramMemoryPreference MEM_PREFERENCE = MemoryPreference; + + /// Whether to dequeue tiles from a global work queue + static constexpr bool IS_WORK_STEALING = WorkStealing; + + /// Vector size for samples loading (1, 2, 4) + static constexpr int VEC_SIZE = VecSize; + static_assert(VEC_SIZE == 1 || VEC_SIZE == 2 || VEC_SIZE == 4); + + ///< The BlockLoad algorithm to use + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + + ///< Cache load modifier for reading input elements + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template +using AgentHistogramPolicy + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceHistogram") = detail::agent_histogram_policy< + ThreadsPerBlock, + PixelsPerThread, + LoadAlgorithm, + LoadModifier, + RleCompress, + MemoryPreference, + WorkStealing, + VecSize>; + +namespace detail::histogram +{ +// Return a native pixel pointer (specialized for CacheModifiedInputIterator types) +template +_CCCL_DEVICE _CCCL_FORCEINLINE auto NativePointer(CacheModifiedInputIterator itr) +{ + return itr.ptr; +} + +// Return a native pixel pointer (specialized for other types) +template +_CCCL_DEVICE _CCCL_FORCEINLINE auto NativePointer(IteratorT itr) +{ + return nullptr; +} + +//! @brief AgentHistogram implements a stateful abstraction of CUDA thread blocks for participating +//! in device-wide histogram . +//! +//! @tparam AgentHistogramPolicyT +//! Parameterized AgentHistogramPolicy tuning policy type +//! +//! @tparam PrivatizedSmemBins +//! Number of privatized shared-memory histogram bins of any channel. Zero indicates privatized +//! counters to be maintained in device-accessible memory. +//! +//! @tparam NumChannels +//! Number of channels interleaved in the input data. Supports up to four channels. +//! +//! @tparam NumActiveChannels +//! Number of channels actively being histogrammed +//! +//! @tparam SampleIteratorT +//! Random-access input iterator type for reading samples +//! +//! @tparam CounterT +//! Integer type for counting sample occurrences per histogram bin +//! +//! @tparam PrivatizedDecodeOpT +//! The transform operator type for determining privatized counter indices from samples, one for +//! each channel +//! +//! @tparam OutputDecodeOpT +//! The transform operator type for determining output bin-ids from privatized counter indices, one +//! for each channel +//! +//! @tparam OffsetT +//! Signed integer type for global offsets +template +struct AgentHistogram +{ + static constexpr int vec_size = AgentHistogramPolicyT::VEC_SIZE; + static constexpr int threads_per_block = AgentHistogramPolicyT::BLOCK_THREADS; + static constexpr int pixels_per_thread = AgentHistogramPolicyT::PIXELS_PER_THREAD; + static constexpr int samples_per_thread = pixels_per_thread * NumChannels; + static constexpr int vecs_per_thread = samples_per_thread / vec_size; + static constexpr int tile_pixels = pixels_per_thread * threads_per_block; + static constexpr int tile_samples = samples_per_thread * threads_per_block; + static constexpr bool is_rle_compress = AgentHistogramPolicyT::IS_RLE_COMPRESS; + static constexpr bool is_work_stealing = AgentHistogramPolicyT::IS_WORK_STEALING; + static constexpr CacheLoadModifier load_modifier = AgentHistogramPolicyT::LOAD_MODIFIER; + static constexpr auto mem_preference = + (PrivatizedSmemBins > 0) ? BlockHistogramMemoryPreference{AgentHistogramPolicyT::MEM_PREFERENCE} : GMEM; + + using SampleT = it_value_t; + using PixelT = typename CubVector::Type; + using VecT = typename CubVector::Type; + + /// Input iterator wrapper type (for applying cache modifier) + // Wrap the native input pointer with CacheModifiedInputIterator or directly use the supplied input iterator type + // TODO(bgruber): we can wrap all contiguous iterators, not just pointers + using WrappedSampleIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + SampleIteratorT>; + using WrappedPixelIteratorT = CacheModifiedInputIterator; + using WrappedVecsIteratorT = CacheModifiedInputIterator; + using BlockLoadSampleT = + BlockLoad; + using BlockLoadPixelT = + BlockLoad; + using BlockLoadVecT = BlockLoad; + + struct _TempStorage + { + // Smem needed for block-privatized smem histogram (with 1 word of padding) + CounterT histograms[NumActiveChannels][PrivatizedSmemBins + 1]; + int tile_idx; + + union + { + typename BlockLoadSampleT::TempStorage sample_load; + typename BlockLoadPixelT::TempStorage pixel_load; + typename BlockLoadVecT::TempStorage vec_load; + }; + }; + + using TempStorage = Uninitialized<_TempStorage>; + + _TempStorage& temp_storage; + WrappedSampleIteratorT d_wrapped_samples; // with cache modifier applied, if possible + SampleT* d_native_samples; // possibly nullptr if unavailable + const int* num_output_bins; // one for each channel + const int* num_privatized_bins; // one for each channel + CounterT* d_privatized_histograms[NumActiveChannels]; // one for each channel + CounterT** d_output_histograms; // in global memory + const OutputDecodeOpT* output_decode_op; // determines output bin-id from privatized counter index, one for each + // channel + const PrivatizedDecodeOpT* privatized_decode_op; // determines privatized counter index from sample, one for each + // channel + bool prefer_smem; // for privatized counterss + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ZeroBinCounters(TwoDimSubscriptableCounterT& privatized_histograms) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + for (int bin = static_cast(threadIdx.x); bin < num_privatized_bins[ch]; bin += threads_per_block) + { + privatized_histograms[ch][bin] = 0; + } + } + + // TODO(bgruber): do we also need the __syncthreads() when prefer_smem is false? + // Barrier to make sure all threads are done updating counters + __syncthreads(); + } + + // Update final output histograms from privatized histograms + template + _CCCL_DEVICE _CCCL_FORCEINLINE void StoreOutput(TwoDimSubscriptableCounterT& privatized_histograms) + { + // Barrier to make sure all threads are done updating counters + __syncthreads(); + + // Apply privatized bin counts to output bin counts + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + const int channel_bins = num_privatized_bins[ch]; + for (int bin = static_cast(threadIdx.x); bin < channel_bins; bin += threads_per_block) + { + int output_bin = -1; + const CounterT count = privatized_histograms[ch][bin]; + const bool is_valid = count > 0; + output_decode_op[ch].template BinSelect(static_cast(bin), output_bin, is_valid); + + if (output_bin >= 0) + { + atomicAdd(&d_output_histograms[ch][output_bin], count); + } + } + } + } + + // Accumulate pixels. Specialized for RLE compression. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void AccumulatePixels( + SampleT samples[pixels_per_thread][NumChannels], + bool is_valid[pixels_per_thread], + TwoDimSubscriptableCounterT& privatized_histograms, + ::cuda::std::true_type is_rle_compress) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + // Bin pixels + int bins[pixels_per_thread]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int pixel = 0; pixel < pixels_per_thread; ++pixel) + { + bins[pixel] = -1; + privatized_decode_op[ch].template BinSelect(samples[pixel][ch], bins[pixel], is_valid[pixel]); + } + + CounterT accumulator = 1; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int pixel = 0; pixel < pixels_per_thread - 1; ++pixel) + { + if (bins[pixel] != bins[pixel + 1]) + { + if (bins[pixel] >= 0) + { + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_60, + (atomicAdd_block(privatized_histograms[ch] + bins[pixel], accumulator);), + (atomicAdd(privatized_histograms[ch] + bins[pixel], accumulator);)); + } + + accumulator = 0; + } + accumulator++; + } + + // Last pixel + if (bins[pixels_per_thread - 1] >= 0) + { + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_60, + (atomicAdd_block(privatized_histograms[ch] + bins[pixels_per_thread - 1], accumulator);), + (atomicAdd(privatized_histograms[ch] + bins[pixels_per_thread - 1], accumulator);)); + } + } + } + + // Accumulate pixels. Specialized for individual accumulation of each pixel. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void AccumulatePixels( + SampleT samples[pixels_per_thread][NumChannels], + bool is_valid[pixels_per_thread], + TwoDimSubscriptableCounterT& privatized_histograms, + ::cuda::std::false_type is_rle_compress) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int pixel = 0; pixel < pixels_per_thread; ++pixel) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + int bin = -1; + privatized_decode_op[ch].template BinSelect(samples[pixel][ch], bin, is_valid[pixel]); + if (bin >= 0) + { + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_60, + (atomicAdd_block(privatized_histograms[ch] + bin, 1);), + (atomicAdd(privatized_histograms[ch] + bin, 1);)); + } + } + } + } + + // Load full, aligned tile using pixel iterator + _CCCL_DEVICE _CCCL_FORCEINLINE void + LoadFullAlignedTile(OffsetT block_offset, SampleT (&samples)[pixels_per_thread][NumChannels]) + { + if constexpr (NumActiveChannels == 1) + { + using AliasedVecs = VecT[vecs_per_thread]; + WrappedVecsIteratorT d_wrapped_vecs(reinterpret_cast(d_native_samples + block_offset)); + // Load using a wrapped vec iterator + BlockLoadVecT{temp_storage.vec_load}.Load(d_wrapped_vecs, reinterpret_cast(samples)); + } + else + { + using AliasedPixels = PixelT[pixels_per_thread]; + WrappedPixelIteratorT d_wrapped_pixels(reinterpret_cast(d_native_samples + block_offset)); + // Load using a wrapped pixel iterator + BlockLoadPixelT{temp_storage.pixel_load}.Load(d_wrapped_pixels, reinterpret_cast(samples)); + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + LoadTile(OffsetT block_offset, int valid_samples, SampleT (&samples)[pixels_per_thread][NumChannels]) + { + if constexpr (IsFullTile) + { + if constexpr (IsAligned) + { + LoadFullAlignedTile(block_offset, samples); + } + else + { + // Load using sample iterator + using AliasedSamples = SampleT[samples_per_thread]; + BlockLoadSampleT{temp_storage.sample_load}.Load( + d_wrapped_samples + block_offset, reinterpret_cast(samples)); + } + } + else + { + if constexpr (IsAligned) + { + // Load partially-full, aligned tile using the pixel iterator + using AliasedPixels = PixelT[pixels_per_thread]; + WrappedPixelIteratorT d_wrapped_pixels((PixelT*) (d_native_samples + block_offset)); + int valid_pixels = valid_samples / NumChannels; + + // Load using a wrapped pixel iterator + BlockLoadPixelT{temp_storage.pixel_load}.Load( + d_wrapped_pixels, reinterpret_cast(samples), valid_pixels); + } + else + { + using AliasedSamples = SampleT[samples_per_thread]; + BlockLoadSampleT{temp_storage.sample_load}.Load( + d_wrapped_samples + block_offset, reinterpret_cast(samples), valid_samples); + } + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void MarkValid(bool (&is_valid)[pixels_per_thread], int valid_samples) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int pixel = 0; pixel < pixels_per_thread; ++pixel) + { + if constexpr (IsStriped) + { + is_valid[pixel] = IsFullTile || (((threadIdx.x + threads_per_block * pixel) * NumChannels) < valid_samples); + } + else + { + is_valid[pixel] = IsFullTile || (((threadIdx.x * pixels_per_thread + pixel) * NumChannels) < valid_samples); + } + } + } + + //! @brief Consume a tile of data samples + //! + //! @tparam IsAligned + //! Whether the tile offset is aligned (vec-aligned for single-channel, pixel-aligned for multi-channel) + //! + //! @tparam IsFullTile + //! Whether the tile is full + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTile(OffsetT block_offset, int valid_samples) + { + SampleT samples[pixels_per_thread][NumChannels]; + bool is_valid[pixels_per_thread]; + + LoadTile(block_offset, valid_samples, samples); + MarkValid(is_valid, valid_samples); + + if (prefer_smem) + { + AccumulatePixels(samples, is_valid, temp_storage.histograms, ::cuda::std::bool_constant{}); + } + else + { + AccumulatePixels(samples, is_valid, d_privatized_histograms, ::cuda::std::bool_constant{}); + } + } + + //! @brief Consume row tiles. Specialized for work-stealing from queue + //! + //! @param num_row_pixels + //! The number of multi-channel pixels per row in the region of interest + //! + //! @param num_rows + //! The number of rows in the region of interest + //! + //! @param row_stride_samples + //! The number of samples between starts of consecutive rows in the region of interest + //! + //! @param tiles_per_row + //! Number of image tiles per row + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTiles( + OffsetT num_row_pixels, + OffsetT num_rows, + OffsetT row_stride_samples, + int tiles_per_row, + GridQueue tile_queue, + ::cuda::std::true_type is_work_stealing) + { + int num_tiles = num_rows * tiles_per_row; + int tile_idx = static_cast((blockIdx.y * gridDim.x) + blockIdx.x); + OffsetT num_even_share_tiles = gridDim.x * gridDim.y; + + while (tile_idx < num_tiles) + { + int row = tile_idx / tiles_per_row; + int col = tile_idx - (row * tiles_per_row); + OffsetT row_offset = row * row_stride_samples; + OffsetT col_offset = (col * tile_samples); + OffsetT tile_offset = row_offset + col_offset; + + if (col == tiles_per_row - 1) + { + // Consume a partially-full tile at the end of the row + OffsetT num_remaining = (num_row_pixels * NumChannels) - col_offset; + ConsumeTile(tile_offset, num_remaining); + } + else + { + // Consume full tile + ConsumeTile(tile_offset, tile_samples); + } + + __syncthreads(); + + // Get next tile + if (threadIdx.x == 0) + { + temp_storage.tile_idx = tile_queue.Drain(1) + num_even_share_tiles; + } + + __syncthreads(); + + tile_idx = temp_storage.tile_idx; + } + } + + //! @brief Consume row tiles. Specialized for even-share (striped across thread blocks) + //! + //! @param num_row_pixels + //! The number of multi-channel pixels per row in the region of interest + //! + //! @param num_rows + //! The number of rows in the region of interest + //! + //! @param row_stride_samples + //! The number of samples between starts of consecutive rows in the region of interest + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTiles( + OffsetT num_row_pixels, OffsetT num_rows, OffsetT row_stride_samples, int, GridQueue, ::cuda::std::false_type) + { + for (int row = static_cast(blockIdx.y); row < num_rows; row += static_cast(gridDim.y)) + { + OffsetT row_begin = row * row_stride_samples; + OffsetT row_end = row_begin + (num_row_pixels * NumChannels); + OffsetT tile_offset = row_begin + (blockIdx.x * tile_samples); + + while (tile_offset < row_end) + { + OffsetT num_remaining = row_end - tile_offset; + + if (num_remaining < tile_samples) + { + // Consume partial tile + ConsumeTile(tile_offset, num_remaining); + break; + } + + // Consume full tile + ConsumeTile(tile_offset, tile_samples); + tile_offset += gridDim.x * tile_samples; + } + } + } + + //--------------------------------------------------------------------- + // Parameter extraction + //--------------------------------------------------------------------- + + //! @brief Constructor + //! + //! @param temp_storage + //! Reference to temp_storage + //! + //! @param d_samples + //! Input data to reduce + //! + //! @param num_output_bins + //! The number bins per final output histogram + //! + //! @param num_privatized_bins + //! The number bins per privatized histogram + //! + //! @param d_output_histograms + //! Reference to final output histograms + //! + //! @param d_privatized_histograms + //! Reference to privatized histograms + //! + //! @param output_decode_op + //! The transform operator for determining output bin-ids from privatized counter indices, one for each channel + //! + //! @param privatized_decode_op + //! The transform operator for determining privatized counter indices from samples, one for each channel + _CCCL_DEVICE _CCCL_FORCEINLINE AgentHistogram( + TempStorage& temp_storage, + SampleIteratorT d_samples, + const int* num_output_bins, + const int* num_privatized_bins, + CounterT** d_output_histograms, + CounterT** d_privatized_histograms, + const OutputDecodeOpT* output_decode_op, + const PrivatizedDecodeOpT* privatized_decode_op) + : temp_storage(temp_storage.Alias()) + , d_wrapped_samples(d_samples) + , d_native_samples(NativePointer(d_wrapped_samples)) + , num_output_bins(num_output_bins) + , num_privatized_bins(num_privatized_bins) + , d_output_histograms(d_output_histograms) + , output_decode_op(output_decode_op) + , privatized_decode_op(privatized_decode_op) + , prefer_smem((mem_preference == SMEM) ? true : // prefer smem privatized histograms + (mem_preference == GMEM) ? false + : // prefer gmem privatized histograms + blockIdx.x & 1) // prefer blended privatized histograms + { + const int blockId = static_cast((blockIdx.y * gridDim.x) + blockIdx.x); + + // TODO(bgruber): d_privatized_histograms seems only used when !prefer_smem, can we skip it if prefer_smem? + // Initialize the locations of this block's privatized histograms + for (int ch = 0; ch < NumActiveChannels; ++ch) + { + const auto offset = static_cast<::cuda::std::int64_t>(blockId) * num_privatized_bins[ch]; + this->d_privatized_histograms[ch] = d_privatized_histograms[ch] + offset; + } + } + + //! @brief Consume image + //! + //! @param num_row_pixels + //! The number of multi-channel pixels per row in the region of interest + //! + //! @param num_rows + //! The number of rows in the region of interest + //! + //! @param row_stride_samples + //! The number of samples between starts of consecutive rows in the region of interest + //! + //! @param tiles_per_row + //! Number of image tiles per row + //! + //! @param tile_queue + //! Queue descriptor for assigning tiles of work to thread blocks + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeTiles( + OffsetT num_row_pixels, OffsetT num_rows, OffsetT row_stride_samples, int tiles_per_row, GridQueue tile_queue) + { + // Check whether all row starting offsets are vec-aligned (in single-channel) or pixel-aligned (in multi-channel) + constexpr int vec_mask = alignof(VecT) - 1; + constexpr int pixel_mask = alignof(PixelT) - 1; + const size_t row_bytes = sizeof(SampleT) * row_stride_samples; + + const bool vec_aligned_rows = + (NumChannels == 1) && (samples_per_thread % vec_size == 0) && // Single channel + ((size_t(d_native_samples) & vec_mask) == 0) && // ptr is quad-aligned + ((num_rows == 1) || ((row_bytes & vec_mask) == 0)); // number of row-samples is a multiple of the alignment of the + // quad + + const bool pixel_aligned_rows = + (NumChannels > 1) && // Multi channel + ((size_t(d_native_samples) & pixel_mask) == 0) && // ptr is pixel-aligned + ((row_bytes & pixel_mask) == 0); // number of row-samples is a multiple of the alignment of the pixel + + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + + // Whether rows are aligned and can be vectorized + if ((d_native_samples != nullptr) && (vec_aligned_rows || pixel_aligned_rows)) + { + ConsumeTiles( + num_row_pixels, num_rows, row_stride_samples, tiles_per_row, tile_queue, bool_constant_v); + } + else + { + ConsumeTiles( + num_row_pixels, num_rows, row_stride_samples, tiles_per_row, tile_queue, bool_constant_v); + } + + _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); // omitting makes no difference in cub.bench.histogram.even.base + } + + //! Initialize privatized bin counters. Specialized for privatized shared-memory counters + _CCCL_DEVICE _CCCL_FORCEINLINE void InitBinCounters() + { + if (prefer_smem) + { + ZeroBinCounters(temp_storage.histograms); + } + else + { + ZeroBinCounters(d_privatized_histograms); + } + } + + //! Store privatized histogram to device-accessible memory. Specialized for privatized shared-memory counters + _CCCL_DEVICE _CCCL_FORCEINLINE void StoreOutput() + { + if (prefer_smem) + { + StoreOutput(temp_storage.histograms); + } + else + { + StoreOutput(d_privatized_histograms); + } + } +}; +} // namespace detail::histogram + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_merge.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_merge.cuh new file mode 100644 index 00000000..d87efe82 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_merge.cuh @@ -0,0 +1,335 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail::merge +{ +// TODO(bgruber): can we unify this one with AgentMerge in agent_merge_sort.cuh? +// TODO(bgruber): pass a merge_policy by value instead of individual template parameters in C++20 +template +struct agent_t +{ + static constexpr int threads_per_block = ThreadsPerBlock; // also used for kernel launch bounds and dispatch logic + static constexpr int items_per_tile = ItemsPerThread * ThreadsPerBlock; // also used by dispatch logic + + // key and value type are taken from the first input sequence (consistent with old Thrust behavior) + using key_type = it_value_t; + using item_type = it_value_t; + + using block_load_to_shared = BlockLoadToShared; + using block_store_keys = BlockStore; + using block_store_items = BlockStore; + + static constexpr int bl2sh_minimum_align = cub::detail::LoadToSharedBufferAlignBytes(); + + template + struct alignas(cub::detail::LoadToSharedBufferAlignBytes()) buffer_t + { + // Need extra bytes of padding for TMA because this static buffer has to hold the two dynamically sized buffers. + static constexpr int bytes_needed = cub::detail::LoadToSharedBufferSizeBytes(items_per_tile + 1ULL) + + (alignof(ValueT) < bl2sh_minimum_align ? 2 * bl2sh_minimum_align : 0); + + char c_array[bytes_needed]; + }; + + struct temp_storages_without_bl2sh + { + using keys_smem = ::cuda::std::conditional_t, key_type[items_per_tile + 1]>; + using items_smem = ::cuda::std::conditional_t, item_type[items_per_tile + 1]>; + union + { + typename block_store_keys::TempStorage store_keys; + typename block_store_items::TempStorage store_items; + keys_smem keys_shared; + items_smem items_shared; + }; + }; + + // inherit from data storage, so it's positioned at the start of the shared memory + struct temp_storages_with_bl2sh : temp_storages_without_bl2sh + { + typename block_load_to_shared::TempStorage load2sh; + }; + + using temp_storages = ::cuda::std:: + conditional_t; + + using TempStorage = Uninitialized; + + // Per thread data + temp_storages& storage; + KeysIt1 keys1_in; + ItemsIt1 items1_in; + Offset keys1_count; + KeysIt2 keys2_in; + ItemsIt2 items2_in; + Offset keys2_count; + KeysOutputIt keys_out; + ItemsOutputIt items_out; + CompareOp compare_op; + Offset* key1_beg_offsets; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void consume_tile(Offset tile_idx, Offset tile_base, int num_remaining) + { + const Offset diag0 = items_per_tile * tile_idx; + Offset diag1 = diag0 + items_per_tile; + if constexpr (IsFullTile) + { + _CCCL_ASSERT(diag1 <= keys1_count + keys2_count, ""); + } + else + { + diag1 = keys1_count + keys2_count; + } + + // compute bounding box for keys1 & keys2 + const Offset keys1_beg = key1_beg_offsets[tile_idx + 0]; + const Offset keys1_end = key1_beg_offsets[tile_idx + 1]; + const Offset keys2_beg = diag0 - keys1_beg; + const Offset keys2_end = diag1 - keys1_end; + + // number of keys per tile + const int keys1_count_tile = static_cast(keys1_end - keys1_beg); + const int keys2_count_tile = static_cast(keys2_end - keys2_beg); + if constexpr (IsFullTile) // NOLINT(bugprone-branch-clone) + { + _CCCL_ASSERT(keys1_count_tile + keys2_count_tile == items_per_tile, ""); + } + else + { + _CCCL_ASSERT(keys1_count_tile + keys2_count_tile == num_remaining, ""); + } + + [[maybe_unused]] auto load2sh = [&] { + if constexpr (UseBl2ShForKeys || UseBl2ShForItems) + { + return block_load_to_shared{storage.load2sh}; + } + else + { + return NullType{}; + } + }(); + + key_type keys_loc[ItemsPerThread]; + key_type* keys1_shared; + key_type* keys2_shared; + int keys2_offset; + if constexpr (UseBl2ShForKeys) + { + ::cuda::std::span keys1_src{THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(keys1_in + keys1_beg), + static_cast<::cuda::std::size_t>(keys1_count_tile)}; + ::cuda::std::span keys2_src{THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(keys2_in + keys2_beg), + static_cast<::cuda::std::size_t>(keys2_count_tile)}; + ::cuda::std::span keys_buffers{storage.keys_shared.c_array}; + auto keys1_buffer = keys_buffers.first(cub::detail::LoadToSharedBufferSizeBytes(keys1_count_tile)); + auto keys2_buffer = keys_buffers.last(cub::detail::LoadToSharedBufferSizeBytes(keys2_count_tile)); + _CCCL_ASSERT(keys1_buffer.end() <= keys2_buffer.begin(), + "Keys buffer needs to be appropriately sized (internal)"); + keys1_shared = data(load2sh.CopyAsync(keys1_buffer, keys1_src)); + keys2_shared = data(load2sh.CopyAsync(keys2_buffer, keys2_src)); + auto token = load2sh.Commit(); + // Needed for using keys1_shared as one big buffer including both ranges in SerialMerge + keys2_offset = static_cast(keys2_shared - keys1_shared); + load2sh.Wait(::cuda::std::move(token)); + } + else + { + auto keys1_in_cm = try_make_cache_modified_iterator(keys1_in); + auto keys2_in_cm = try_make_cache_modified_iterator(keys2_in); + merge_sort::gmem_to_reg( + keys_loc, keys1_in_cm + keys1_beg, keys2_in_cm + keys2_beg, keys1_count_tile, keys2_count_tile); + keys1_shared = &storage.keys_shared[0]; + // Needed for using keys1_shared as one big buffer including both ranges in SerialMerge + keys2_offset = keys1_count_tile; + keys2_shared = keys1_shared + keys2_offset; + merge_sort::reg_to_shared(keys1_shared, keys_loc); + __syncthreads(); + } + + // Now find the merge path for each of the threads. + // We can use int type here, because the number of items in shared memory is limited. + int diag0_thread = ItemsPerThread * static_cast(threadIdx.x); + if constexpr (IsFullTile) + { + _CCCL_ASSERT(num_remaining == items_per_tile, ""); + _CCCL_ASSERT(diag0_thread < num_remaining, ""); + } + else + { // for partial tiles, clamp the thread diagonal to the valid items + diag0_thread = (::cuda::std::min) (diag0_thread, num_remaining); + } + + const int keys1_beg_thread = + MergePath(keys1_shared, keys2_shared, keys1_count_tile, keys2_count_tile, diag0_thread, compare_op); + const int keys2_beg_thread = diag0_thread - keys1_beg_thread; + + const int keys1_count_thread = keys1_count_tile - keys1_beg_thread; + const int keys2_count_thread = keys2_count_tile - keys2_beg_thread; + + // perform serial merge + int indices[ItemsPerThread]; + cub::detail::serial_merge( + keys1_shared, + keys1_beg_thread, + keys2_offset + keys2_beg_thread, + keys1_count_thread, + keys2_count_thread, + keys_loc, + indices, + compare_op); + + // write keys + __syncthreads(); // sync after reading from SMEM before so block store can use SMEM again + if constexpr (IsFullTile) + { + block_store_keys{storage.store_keys}.Store(keys_out + tile_base, keys_loc); + } + else + { + block_store_keys{storage.store_keys}.Store(keys_out + tile_base, keys_loc, num_remaining); + } + + // if items are provided, merge them + static constexpr bool have_items = !::cuda::std::is_same_v; + if constexpr (have_items) + { + // Both of these are only needed when either keys or items or both use BlockLoadToShared introducing padding (that + // can differ between the keys and items) + [[maybe_unused]] const auto translate_indices = [&](int items2_offset) -> void { + const int diff = items2_offset - keys2_offset; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; ++i) + { + if (indices[i] >= keys2_offset) + { + indices[i] += diff; + } + } + }; + // WAR for MSVC erroring ("declared but never referenced") despite [[maybe_unused]] + (void) translate_indices; + + item_type items_loc[ItemsPerThread]; + item_type* items1_shared; + if constexpr (UseBl2ShForItems) + { + ::cuda::std::span items1_src{THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(items1_in + keys1_beg), + static_cast<::cuda::std::size_t>(keys1_count_tile)}; + ::cuda::std::span items2_src{THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(items2_in + keys2_beg), + static_cast<::cuda::std::size_t>(keys2_count_tile)}; + ::cuda::std::span items_buffers{storage.items_shared.c_array}; + auto items1_buffer = items_buffers.first(cub::detail::LoadToSharedBufferSizeBytes(keys1_count_tile)); + auto items2_buffer = items_buffers.last(cub::detail::LoadToSharedBufferSizeBytes(keys2_count_tile)); + _CCCL_ASSERT(items1_buffer.end() <= items2_buffer.begin(), + "Items buffer needs to be appropriately sized (internal)"); + // block_store_keys above uses shared memory, so make sure all threads are done before we write + __syncthreads(); + items1_shared = data(load2sh.CopyAsync(items1_buffer, items1_src)); + item_type* items2_shared = data(load2sh.CopyAsync(items2_buffer, items2_src)); + auto token = load2sh.Commit(); + const int items2_offset = static_cast(items2_shared - items1_shared); + translate_indices(items2_offset); + load2sh.Wait(::cuda::std::move(token)); + } + else + { + { + auto items1_in_cm = try_make_cache_modified_iterator(items1_in); + auto items2_in_cm = try_make_cache_modified_iterator(items2_in); + merge_sort::gmem_to_reg( + items_loc, items1_in_cm + keys1_beg, items2_in_cm + keys2_beg, keys1_count_tile, keys2_count_tile); + __syncthreads(); // block_store_keys above uses SMEM, so make sure all threads are done before we write to it + items1_shared = &storage.items_shared[0]; + if constexpr (UseBl2ShForKeys) + { + const int items2_offset = keys1_count_tile; + translate_indices(items2_offset); + } + merge_sort::reg_to_shared(items1_shared, items_loc); + __syncthreads(); + } + } + + // gather items from shared mem + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; ++i) + { + items_loc[i] = items1_shared[indices[i]]; + } + __syncthreads(); + + // write from reg to gmem + if constexpr (IsFullTile) + { + block_store_items{storage.store_items}.Store(items_out + tile_base, items_loc); + } + else + { + block_store_items{storage.store_items}.Store(items_out + tile_base, items_loc, num_remaining); + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + const Offset tile_idx = blockIdx.x; + const Offset tile_base = tile_idx * items_per_tile; + const int items_in_tile = + static_cast((::cuda::std::min) (static_cast(items_per_tile), keys1_count + keys2_count - tile_base)); + if (items_in_tile == items_per_tile) + { + consume_tile(tile_idx, tile_base, items_per_tile); + } + else + { + consume_tile(tile_idx, tile_base, items_in_tile); + } + } +}; +} // namespace detail::merge +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_merge_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_merge_sort.cuh new file mode 100644 index 00000000..929ea52d --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_merge_sort.cuh @@ -0,0 +1,692 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN +namespace detail::merge_sort +{ +template +struct AgentBlockSort +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + + static constexpr MergeSortPolicy policy = PolicyGetter{}(); + static constexpr int BLOCK_THREADS = policy.threads_per_block; + static constexpr int ITEMS_PER_THREAD = policy.items_per_thread; + static constexpr int ITEMS_PER_TILE = BLOCK_THREADS * ITEMS_PER_THREAD; + + using BlockMergeSortT = BlockMergeSort; + + using KeysLoadIt = try_make_cache_modified_iterator_t; + using ItemsLoadIt = try_make_cache_modified_iterator_t; + + using BlockLoadKeys = BlockLoad, BLOCK_THREADS, ITEMS_PER_THREAD, policy.load_algorithm>; + using BlockLoadItems = BlockLoad, BLOCK_THREADS, ITEMS_PER_THREAD, policy.load_algorithm>; + + using BlockStoreKeysIt = + BlockStore, BLOCK_THREADS, ITEMS_PER_THREAD, policy.store_algorithm>; + using BlockStoreItemsIt = + BlockStore, BLOCK_THREADS, ITEMS_PER_THREAD, policy.store_algorithm>; + using BlockStoreKeysRaw = BlockStore; + using BlockStoreItemsRaw = BlockStore; + + union _TempStorage + { + typename BlockLoadKeys::TempStorage load_keys; + typename BlockLoadItems::TempStorage load_items; + typename BlockStoreKeysIt::TempStorage store_keys_it; + typename BlockStoreItemsIt::TempStorage store_items_it; + typename BlockStoreKeysRaw::TempStorage store_keys_raw; + typename BlockStoreItemsRaw::TempStorage store_items_raw; + typename BlockMergeSortT::TempStorage block_merge; + }; + + /// Alias wrapper allowing storage to be unioned + using TempStorage = Uninitialized<_TempStorage>; + + //--------------------------------------------------------------------- + // Per thread data + //--------------------------------------------------------------------- + + bool ping; + _TempStorage& storage; + KeysLoadIt keys_in; + ItemsLoadIt items_in; + OffsetT keys_count; + KeyIteratorT keys_out_it; + ValueIteratorT items_out_it; + KeyT* keys_out_raw; + ValueT* items_out_raw; + CompareOpT compare_op; + + _CCCL_DEVICE _CCCL_FORCEINLINE AgentBlockSort( + bool ping_, + TempStorage& storage_, + KeysLoadIt keys_in_, + ItemsLoadIt items_in_, + OffsetT keys_count_, + KeyIteratorT keys_out_it_, + ValueIteratorT items_out_it_, + KeyT* keys_out_raw_, + ValueT* items_out_raw_, + CompareOpT compare_op_) + : ping(ping_) + , storage(storage_.Alias()) + , keys_in(keys_in_) + , items_in(items_in_) + , keys_count(keys_count_) + , keys_out_it(keys_out_it_) + , items_out_it(items_out_it_) + , keys_out_raw(keys_out_raw_) + , items_out_raw(items_out_raw_) + , compare_op(compare_op_) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE void Process() + { + const auto tile_idx = static_cast(blockIdx.x); + const auto num_tiles = static_cast(gridDim.x); + const auto tile_base = tile_idx * ITEMS_PER_TILE; + const int items_in_tile = (::cuda::std::min) (static_cast(keys_count - tile_base), int{ITEMS_PER_TILE}); + + if (tile_idx < num_tiles - 1) + { + consume_tile(tile_base, ITEMS_PER_TILE); + } + else + { + consume_tile(tile_base, items_in_tile); + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void consume_tile(OffsetT tile_base, int num_remaining) + { + ValueT items_local[ITEMS_PER_THREAD]; + + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + + if constexpr (!KEYS_ONLY) + { + if constexpr (IS_LAST_TILE) + { + BlockLoadItems(storage.load_items) + .Load(items_in + tile_base, items_local, num_remaining, *(items_in + tile_base)); + } + else + { + BlockLoadItems(storage.load_items).Load(items_in + tile_base, items_local); + } + + __syncthreads(); + } + + KeyT keys_local[ITEMS_PER_THREAD]; + if constexpr (IS_LAST_TILE) + { + BlockLoadKeys(storage.load_keys).Load(keys_in + tile_base, keys_local, num_remaining, *(keys_in + tile_base)); + } + else + { + BlockLoadKeys(storage.load_keys).Load(keys_in + tile_base, keys_local); + } + + __syncthreads(); + _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); + + if constexpr (IS_LAST_TILE) + { + BlockMergeSortT(storage.block_merge).Sort(keys_local, items_local, compare_op, num_remaining, keys_local[0]); + } + else + { + BlockMergeSortT(storage.block_merge).Sort(keys_local, items_local, compare_op); + } + + __syncthreads(); + + if (ping) + { + if constexpr (IS_LAST_TILE) + { + BlockStoreKeysIt(storage.store_keys_it).Store(keys_out_it + tile_base, keys_local, num_remaining); + } + else + { + BlockStoreKeysIt(storage.store_keys_it).Store(keys_out_it + tile_base, keys_local); + } + + if constexpr (!KEYS_ONLY) + { + __syncthreads(); + + if constexpr (IS_LAST_TILE) + { + BlockStoreItemsIt(storage.store_items_it).Store(items_out_it + tile_base, items_local, num_remaining); + } + else + { + BlockStoreItemsIt(storage.store_items_it).Store(items_out_it + tile_base, items_local); + } + } + } + else + { + if constexpr (IS_LAST_TILE) + { + BlockStoreKeysRaw(storage.store_keys_raw).Store(keys_out_raw + tile_base, keys_local, num_remaining); + } + else + { + BlockStoreKeysRaw(storage.store_keys_raw).Store(keys_out_raw + tile_base, keys_local); + } + + if constexpr (!KEYS_ONLY) + { + __syncthreads(); + + if constexpr (IS_LAST_TILE) + { + BlockStoreItemsRaw(storage.store_items_raw).Store(items_out_raw + tile_base, items_local, num_remaining); + } + else + { + BlockStoreItemsRaw(storage.store_items_raw).Store(items_out_raw + tile_base, items_local); + } + } + } + } +}; + +/** + * \brief This agent is responsible for partitioning a merge path into equal segments + * + * There are two sorted arrays to be merged into one array. If the first array + * is partitioned between parallel workers by slicing it into ranges of equal + * size, there could be a significant workload imbalance. The imbalance is + * caused by the fact that the distribution of elements from the second array + * is unknown beforehand. Instead, the MergePath is partitioned between workers. + * This approach guarantees an equal amount of work being assigned to each worker. + * + * This approach is outlined in the paper: + * Odeh et al, "Merge Path - Parallel Merging Made Simple" + * doi:10.1109/IPDPSW.2012.202 + */ +template +struct AgentPartition +{ + bool ping; + KeyIteratorT keys_ping; + KeyT* keys_pong; + OffsetT keys_count; + OffsetT partition_idx; + OffsetT* merge_partitions; + CompareOpT compare_op; + OffsetT target_merged_tiles_number; + int items_per_tile; + OffsetT num_partitions; + + _CCCL_DEVICE _CCCL_FORCEINLINE void Process() + { + const OffsetT merged_tiles_number = target_merged_tiles_number / 2; + + // target_merged_tiles_number is a power of two. + const OffsetT mask = target_merged_tiles_number - 1; + + // The first tile number in the tiles group being merged, equal to: + // target_merged_tiles_number * (partition_idx / target_merged_tiles_number) + const OffsetT list = ~mask & partition_idx; + const OffsetT start = items_per_tile * list; + const OffsetT size = items_per_tile * merged_tiles_number; + + // Tile number within the tile group being merged, equal to: + // partition_idx / target_merged_tiles_number + const OffsetT local_tile_idx = mask & partition_idx; + + const OffsetT keys1_beg = (::cuda::std::min) (keys_count, start); + const OffsetT keys1_end = (::cuda::std::min) (keys_count, detail::safe_add_bound_to_max(start, size)); + const OffsetT keys2_beg = keys1_end; + const OffsetT keys2_end = (::cuda::std::min) (keys_count, detail::safe_add_bound_to_max(keys2_beg, size)); + + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + + // The last partition (which is one-past-the-last-tile) is only to mark the end of keys1_end for the merge stage + if (partition_idx + 1 == num_partitions) + { + merge_partitions[partition_idx] = keys1_end; + } + else + { + const OffsetT partition_at = (::cuda::std::min) (keys2_end - keys1_beg, items_per_tile * local_tile_idx); + + OffsetT partition_diag = + ping + ? MergePath(keys_ping + keys1_beg, + keys_ping + keys2_beg, + keys1_end - keys1_beg, + keys2_end - keys2_beg, + partition_at, + compare_op) + : MergePath(keys_pong + keys1_beg, + keys_pong + keys2_beg, + keys1_end - keys1_beg, + keys2_end - keys2_beg, + partition_at, + compare_op); + + merge_partitions[partition_idx] = keys1_beg + partition_diag; + } + + // TODO(bgruber): looking at SASS triggering the next launch here just generates a lot of noise and the PRE-EXIT + // just ends of right before EXIT anyway. So let's omit it. + // _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); + } +}; + +/** + * \brief Concatenates up to ITEMS_PER_THREAD elements from input{1,2} into output array + * + * Reads data in a coalesced fashion [BLOCK_THREADS * item + tid] and + * stores the result in output[item]. + */ +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +gmem_to_reg(T (&output)[ITEMS_PER_THREAD], It1 input1, It2 input2, int count1, int count2) +{ + if constexpr (IS_FULL_TILE) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ITEMS_PER_THREAD; ++item) + { + const int idx = BLOCK_THREADS * item + threadIdx.x; + // It1 and It2 could have different value types. Convert after load. + output[item] = (idx < count1) ? static_cast(input1[idx]) : static_cast(input2[idx - count1]); + } + } + else + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ITEMS_PER_THREAD; ++item) + { + const int idx = BLOCK_THREADS * item + threadIdx.x; + if (idx < count1 + count2) + { + output[item] = (idx < count1) ? static_cast(input1[idx]) : static_cast(input2[idx - count1]); + } + } + } +} + +/// \brief Stores data in a coalesced fashion in[item] -> out[BLOCK_THREADS * item + tid] +template +_CCCL_DEVICE _CCCL_FORCEINLINE void reg_to_shared(It output, T (&input)[ITEMS_PER_THREAD]) +{ + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ITEMS_PER_THREAD; ++item) + { + const int idx = BLOCK_THREADS * item + threadIdx.x; + output[idx] = input[item]; + } +} + +/// \brief The agent is responsible for merging N consecutive sorted arrays into N/2 sorted arrays. +template +struct AgentMerge +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + + static constexpr MergeSortPolicy policy = PolicyGetter{}(); + static constexpr int BLOCK_THREADS = policy.threads_per_block; + static constexpr int ITEMS_PER_THREAD = policy.items_per_thread; + static constexpr int ITEMS_PER_TILE = BLOCK_THREADS * ITEMS_PER_THREAD; + + using KeysLoadPingIt = try_make_cache_modified_iterator_t; + using ItemsLoadPingIt = try_make_cache_modified_iterator_t; + using KeysLoadPongIt = try_make_cache_modified_iterator_t; + using ItemsLoadPongIt = try_make_cache_modified_iterator_t; + + using KeysOutputPongIt = KeyIteratorT; + using ItemsOutputPongIt = ValueIteratorT; + using KeysOutputPingIt = KeyT*; + using ItemsOutputPingIt = ValueT*; + + using BlockStoreKeysPong = + BlockStore, BLOCK_THREADS, ITEMS_PER_THREAD, policy.store_algorithm>; + using BlockStoreItemsPong = + BlockStore, BLOCK_THREADS, ITEMS_PER_THREAD, policy.store_algorithm>; + + using BlockStoreKeysPing = + BlockStore, BLOCK_THREADS, ITEMS_PER_THREAD, policy.store_algorithm>; + using BlockStoreItemsPing = + BlockStore, BLOCK_THREADS, ITEMS_PER_THREAD, policy.store_algorithm>; + + /// Parameterized BlockReduce primitive + + union _TempStorage + { + typename BlockStoreKeysPing::TempStorage store_keys_ping; + typename BlockStoreItemsPing::TempStorage store_items_ping; + typename BlockStoreKeysPong::TempStorage store_keys_pong; + typename BlockStoreItemsPong::TempStorage store_items_pong; + + KeyT keys_shared[ITEMS_PER_TILE + 1]; + ValueT items_shared[ITEMS_PER_TILE + 1]; + }; + + /// Alias wrapper allowing storage to be unioned + using TempStorage = Uninitialized<_TempStorage>; + + //--------------------------------------------------------------------- + // Per thread data + //--------------------------------------------------------------------- + + bool ping; + _TempStorage& storage; + + KeysLoadPingIt keys_in_ping; + ItemsLoadPingIt items_in_ping; + KeysLoadPongIt keys_in_pong; + ItemsLoadPongIt items_in_pong; + + OffsetT keys_count; + + KeysOutputPongIt keys_out_pong; + ItemsOutputPongIt items_out_pong; + KeysOutputPingIt keys_out_ping; + ItemsOutputPingIt items_out_ping; + + CompareOpT compare_op; + OffsetT* merge_partitions; + OffsetT target_merged_tiles_number; + + //--------------------------------------------------------------------- + // Utility functions + //--------------------------------------------------------------------- + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void consume_tile(int tid, OffsetT tile_idx, OffsetT tile_base, int count) + { + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + + const OffsetT partition_beg = merge_partitions[tile_idx + 0]; + const OffsetT partition_end = merge_partitions[tile_idx + 1]; + + // target_merged_tiles_number is a power of two. + const OffsetT merged_tiles_number = target_merged_tiles_number / 2; + + const OffsetT mask = target_merged_tiles_number - 1; + + // The first tile number in the tiles group being merged, equal to: + // target_merged_tiles_number * (tile_idx / target_merged_tiles_number) + const OffsetT list = ~mask & tile_idx; + const OffsetT start = ITEMS_PER_TILE * list; + const OffsetT size = ITEMS_PER_TILE * merged_tiles_number; + + const OffsetT diag = ITEMS_PER_TILE * tile_idx - start; + + const OffsetT keys1_beg = partition_beg - start; + OffsetT keys1_end = partition_end - start; + + const OffsetT keys_end_dist_from_start = keys_count - start; + const OffsetT max_keys2 = (keys_end_dist_from_start > size) ? (keys_end_dist_from_start - size) : 0; + + // We have the following invariants: + // diag >= keys1_beg, because diag is the distance of the total merge path so far (keys1 + keys2) + // diag+ITEMS_PER_TILE >= keys1_end, because diag+ITEMS_PER_TILE is the distance of the merge path for the next tile + // and keys1_end is key1's component of that path + const OffsetT keys2_beg = (::cuda::std::min) (max_keys2, diag - keys1_beg); + OffsetT keys2_end = + (::cuda::std::min) (max_keys2, + detail::safe_add_bound_to_max(diag, static_cast(ITEMS_PER_TILE)) - keys1_end); + + // Check if it's the last tile in the tile group being merged + if (mask == (mask & tile_idx)) + { + keys1_end = (::cuda::std::min) (keys_count - start, size); + keys2_end = (::cuda::std::min) (max_keys2, size); + } + + // number of keys per tile + const int num_keys1 = static_cast(keys1_end - keys1_beg); + const int num_keys2 = static_cast(keys2_end - keys2_beg); + + // load keys1 & keys2 + KeyT keys_local[ITEMS_PER_THREAD]; + if (ping) + { + gmem_to_reg( + keys_local, keys_in_ping + start + keys1_beg, keys_in_ping + start + size + keys2_beg, num_keys1, num_keys2); + } + else + { + gmem_to_reg( + keys_local, keys_in_pong + start + keys1_beg, keys_in_pong + start + size + keys2_beg, num_keys1, num_keys2); + } + reg_to_shared(&storage.keys_shared[0], keys_local); + + // preload items into registers already + // + [[maybe_unused]] ValueT items_local[ITEMS_PER_THREAD]; + if constexpr (!KEYS_ONLY) + { + if (ping) + { + gmem_to_reg( + items_local, + items_in_ping + start + keys1_beg, + items_in_ping + start + size + keys2_beg, + num_keys1, + num_keys2); + } + else + { + gmem_to_reg( + items_local, + items_in_pong + start + keys1_beg, + items_in_pong + start + size + keys2_beg, + num_keys1, + num_keys2); + } + } + + __syncthreads(); + _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); + + // use binary search in shared memory + // to find merge path for each of thread + // we can use int type here, because the number of + // items in shared memory is limited + // + const int diag0_local = (::cuda::std::min) (num_keys1 + num_keys2, ITEMS_PER_THREAD * tid); + + const int keys1_beg_local = MergePath( + &storage.keys_shared[0], &storage.keys_shared[num_keys1], num_keys1, num_keys2, diag0_local, compare_op); + const int keys1_end_local = num_keys1; + const int keys2_beg_local = diag0_local - keys1_beg_local; + const int keys2_end_local = num_keys2; + + const int num_keys1_local = keys1_end_local - keys1_beg_local; + const int num_keys2_local = keys2_end_local - keys2_beg_local; + + // perform serial merge + // + int indices[ITEMS_PER_THREAD]; + + detail::serial_merge( + &storage.keys_shared[0], + keys1_beg_local, + keys2_beg_local + num_keys1, + num_keys1_local, + num_keys2_local, + keys_local, + indices, + compare_op); + + __syncthreads(); + + // write keys + if (ping) + { + if constexpr (IS_FULL_TILE) + { + BlockStoreKeysPing(storage.store_keys_ping).Store(keys_out_ping + tile_base, keys_local); + } + else + { + BlockStoreKeysPing(storage.store_keys_ping).Store(keys_out_ping + tile_base, keys_local, num_keys1 + num_keys2); + } + } + else + { + if constexpr (IS_FULL_TILE) + { + BlockStoreKeysPong(storage.store_keys_pong).Store(keys_out_pong + tile_base, keys_local); + } + else + { + BlockStoreKeysPong(storage.store_keys_pong).Store(keys_out_pong + tile_base, keys_local, num_keys1 + num_keys2); + } + } + + // if items are provided, merge them + if constexpr (!KEYS_ONLY) + { + __syncthreads(); + + reg_to_shared(&storage.items_shared[0], items_local); + + __syncthreads(); + + // gather items from shared mem + // + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ITEMS_PER_THREAD; ++item) + { + items_local[item] = storage.items_shared[indices[item]]; + } + + __syncthreads(); + + // write from reg to gmem + // + if (ping) + { + if constexpr (IS_FULL_TILE) + { + BlockStoreItemsPing(storage.store_items_ping).Store(items_out_ping + tile_base, items_local); + } + else + { + BlockStoreItemsPing(storage.store_items_ping).Store(items_out_ping + tile_base, items_local, count); + } + } + else + { + if constexpr (IS_FULL_TILE) + { + BlockStoreItemsPong(storage.store_items_pong).Store(items_out_pong + tile_base, items_local); + } + else + { + BlockStoreItemsPong(storage.store_items_pong).Store(items_out_pong + tile_base, items_local, count); + } + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE AgentMerge( + bool ping_, + TempStorage& storage_, + KeysLoadPingIt keys_in_ping_, + ItemsLoadPingIt items_in_ping_, + KeysLoadPongIt keys_in_pong_, + ItemsLoadPongIt items_in_pong_, + OffsetT keys_count_, + KeysOutputPingIt keys_out_ping_, + ItemsOutputPingIt items_out_ping_, + KeysOutputPongIt keys_out_pong_, + ItemsOutputPongIt items_out_pong_, + CompareOpT compare_op_, + OffsetT* merge_partitions_, + OffsetT target_merged_tiles_number_) + : ping(ping_) + , storage(storage_.Alias()) + , keys_in_ping(keys_in_ping_) + , items_in_ping(items_in_ping_) + , keys_in_pong(keys_in_pong_) + , items_in_pong(items_in_pong_) + , keys_count(keys_count_) + , keys_out_pong(keys_out_pong_) + , items_out_pong(items_out_pong_) + , keys_out_ping(keys_out_ping_) + , items_out_ping(items_out_ping_) + , compare_op(compare_op_) + , merge_partitions(merge_partitions_) + , target_merged_tiles_number(target_merged_tiles_number_) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE void Process() + { + const int tile_idx = static_cast(blockIdx.x); + const int num_tiles = static_cast(gridDim.x); + const OffsetT tile_base = OffsetT(tile_idx) * ITEMS_PER_TILE; + const int tid = static_cast(threadIdx.x); + const int items_in_tile = + static_cast((::cuda::std::min) (static_cast(ITEMS_PER_TILE), keys_count - tile_base)); + + if (tile_idx < num_tiles - 1) + { + consume_tile(tid, tile_idx, tile_base, ITEMS_PER_TILE); + } + else + { + consume_tile(tid, tile_idx, tile_base, items_in_tile); + } + } +}; +} // namespace detail::merge_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_downsweep.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_downsweep.cuh new file mode 100644 index 00000000..84008341 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_downsweep.cuh @@ -0,0 +1,758 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * AgentRadixSortDownsweep implements a stateful abstraction of CUDA thread + * blocks for participating in device-wide radix sort downsweep . + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +/** + * @brief Parameterizable tuning policy type for AgentRadixSortDownsweep + * + * @tparam NominalThreadsPerBlock4B + * Threads per thread block + * + * @tparam NominalItemsPerThread4B + * Items per thread (per tile of input) + * + * @tparam ComputeT + * Dominant compute type + * + * @tparam LoadAlgorithm + * The BlockLoad algorithm to use + * + * @tparam LoadModifier + * Cache load modifier for reading keys (and values) + * + * @tparam RankAlgorithm + * The radix ranking algorithm to use + * + * @tparam ScanAlgorithm + * The block scan algorithm to use + * + * @tparam RadixBits + * The number of radix bits, i.e., log2(bins) + */ +template > +struct agent_radix_sort_downsweep_policy : ScalingType +{ + /// The number of radix bits, i.e., log2(bins) + static constexpr int RADIX_BITS = RadixBits; + + /// The BlockLoad algorithm to use + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + + /// Cache load modifier for reading keys (and values) + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; + + /// The radix ranking algorithm to use + static constexpr RadixRankAlgorithm RANK_ALGORITHM = RankAlgorithm; + + /// The BlockScan algorithm to use + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; +}; +} // namespace detail + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +//! Deprecated [Since 3.5] +template > +using AgentRadixSortDownsweepPolicy + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceRadixSort") = detail::agent_radix_sort_downsweep_policy< + NominalThreadsPerBlock4B, + NominalItemsPerThread4B, + ComputeT, + LoadAlgorithm, + LoadModifier, + RankAlgorithm, + ScanAlgorithm, + RadixBits, + ScalingType>; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::radix_sort +{ +/** + * @brief AgentRadixSortDownsweep implements a stateful abstraction of CUDA thread blocks for participating in + * device-wide radix sort downsweep . + * + * @tparam AgentRadixSortDownsweepPolicy + * Parameterized AgentRadixSortDownsweepPolicy tuning policy type + * + * @tparam IS_DESCENDING + * Whether or not the sorted-order is high-to-low + * + * @tparam KeyT + * KeyT type + * + * @tparam ValueT + * ValueT type + * + * @tparam OffsetT + * Signed integer type for global offsets + */ +template +struct AgentRadixSortDownsweep +{ + //--------------------------------------------------------------------- + // Type definitions and constants + //--------------------------------------------------------------------- + + using traits = radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion = typename traits::bit_ordered_conversion_policy; + + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = AgentRadixSortDownsweepPolicy::LOAD_ALGORITHM; + static constexpr CacheLoadModifier LOAD_MODIFIER = AgentRadixSortDownsweepPolicy::LOAD_MODIFIER; + static constexpr RadixRankAlgorithm RANK_ALGORITHM = AgentRadixSortDownsweepPolicy::RANK_ALGORITHM; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = AgentRadixSortDownsweepPolicy::SCAN_ALGORITHM; + + static constexpr int BLOCK_THREADS = AgentRadixSortDownsweepPolicy::BLOCK_THREADS; + static constexpr int ITEMS_PER_THREAD = AgentRadixSortDownsweepPolicy::ITEMS_PER_THREAD; + static constexpr int RADIX_BITS = AgentRadixSortDownsweepPolicy::RADIX_BITS; + static constexpr int TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD; + + static constexpr int RADIX_DIGITS = 1 << RADIX_BITS; + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + static constexpr bool LOAD_WARP_STRIPED = + RANK_ALGORITHM == RADIX_RANK_MATCH || RANK_ALGORITHM == RADIX_RANK_MATCH_EARLY_COUNTS_ANY + || RANK_ALGORITHM == RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR; + + // Input iterator wrapper type (for applying cache modifier)s + using KeysItr = CacheModifiedInputIterator; + using ValuesItr = CacheModifiedInputIterator; + + // Radix ranking type to use + using BlockRadixRankT = block_radix_rank_t; + + // Digit extractor type + using fundamental_digit_extractor_t = BFEDigitExtractor; + using digit_extractor_t = typename traits::template digit_extractor_t; + + /// Number of bin-starting offsets tracked per thread + static constexpr int BINS_TRACKED_PER_THREAD = BlockRadixRankT::BINS_TRACKED_PER_THREAD; + + // BlockLoad type (keys) + using BlockLoadKeysT = BlockLoad; + + // BlockLoad type (values) + using BlockLoadValuesT = BlockLoad; + + // Value exchange array type + using ValueExchangeT = ValueT[TILE_ITEMS]; + + /** + * Shared memory storage layout + */ + union __align__(16) _TempStorage + { + typename BlockLoadKeysT::TempStorage load_keys; + typename BlockLoadValuesT::TempStorage load_values; + typename BlockRadixRankT::TempStorage radix_rank; + + struct KeysAndOffsets + { + bit_ordered_type exchange_keys[TILE_ITEMS]; + OffsetT relative_bin_offsets[RADIX_DIGITS]; + } keys_and_offsets; + + Uninitialized exchange_values; + + OffsetT exclusive_digit_prefix[RADIX_DIGITS]; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Thread fields + //--------------------------------------------------------------------- + + // Shared storage for this CTA + _TempStorage& temp_storage; + + // Input and output device pointers + KeysItr d_keys_in; + ValuesItr d_values_in; + bit_ordered_type* d_keys_out; + ValueT* d_values_out; + + // The global scatter base offset for each digit (valid in the first RADIX_DIGITS threads) + OffsetT bin_offset[BINS_TRACKED_PER_THREAD]; + + uint32_t current_bit; + uint32_t num_bits; + + // Whether to short-circuit + int short_circuit; + + DecomposerT decomposer; + + //--------------------------------------------------------------------- + // Utility methods + //--------------------------------------------------------------------- + + _CCCL_DEVICE _CCCL_FORCEINLINE digit_extractor_t digit_extractor() + { + return traits::template digit_extractor(current_bit, num_bits, decomposer); + } + + /** + * Scatter ranked keys through shared memory, then to device-accessible memory + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterKeys( + bit_ordered_type (&twiddled_keys)[ITEMS_PER_THREAD], + OffsetT (&relative_bin_offsets)[ITEMS_PER_THREAD], + int (&ranks)[ITEMS_PER_THREAD], + OffsetT valid_items) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + temp_storage.keys_and_offsets.exchange_keys[ranks[ITEM]] = twiddled_keys[ITEM]; + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + bit_ordered_type key = temp_storage.keys_and_offsets.exchange_keys[threadIdx.x + (ITEM * BLOCK_THREADS)]; + uint32_t digit = digit_extractor().Digit(key); + relative_bin_offsets[ITEM] = temp_storage.keys_and_offsets.relative_bin_offsets[digit]; + + key = bit_ordered_conversion::from_bit_ordered(decomposer, key); + + if (FULL_TILE + || (static_cast(threadIdx.x + (ITEM * BLOCK_THREADS)) // NOLINT(bugprone-misplaced-widening-cast) + < valid_items)) + { + d_keys_out[relative_bin_offsets[ITEM] + threadIdx.x + (ITEM * BLOCK_THREADS)] = key; + } + } + } + + /** + * Scatter ranked values through shared memory, then to device-accessible memory + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterValues( + ValueT (&values)[ITEMS_PER_THREAD], + OffsetT (&relative_bin_offsets)[ITEMS_PER_THREAD], + int (&ranks)[ITEMS_PER_THREAD], + OffsetT valid_items) + { + __syncthreads(); + + ValueExchangeT& exchange_values = temp_storage.exchange_values.Alias(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + exchange_values[ranks[ITEM]] = values[ITEM]; + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + ValueT value = exchange_values[threadIdx.x + (ITEM * BLOCK_THREADS)]; + + if (FULL_TILE + || (static_cast(threadIdx.x + (ITEM * BLOCK_THREADS)) // NOLINT(bugprone-misplaced-widening-cast) + < valid_items)) + { + d_values_out[relative_bin_offsets[ITEM] + threadIdx.x + (ITEM * BLOCK_THREADS)] = value; + } + } + } + + /** + * Load a tile of keys (specialized for full tile, block load) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadKeys( + bit_ordered_type (&keys)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + bit_ordered_type oob_item, + ::cuda::std::true_type is_full_tile, + ::cuda::std::false_type warp_striped) + { + BlockLoadKeysT(temp_storage.load_keys).Load(d_keys_in + block_offset, keys); + + __syncthreads(); + } + + /** + * Load a tile of keys (specialized for partial tile, block load) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadKeys( + bit_ordered_type (&keys)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + bit_ordered_type oob_item, + ::cuda::std::false_type is_full_tile, + ::cuda::std::false_type warp_striped) + { + // Register pressure work-around: moving valid_items through shfl prevents compiler + // from reusing guards/addressing from prior guarded loads + valid_items = ::cuda::device::warp_shuffle_idx(valid_items, 0); + + BlockLoadKeysT(temp_storage.load_keys).Load(d_keys_in + block_offset, keys, valid_items, oob_item); + + __syncthreads(); + } + + /** + * Load a tile of keys (specialized for full tile, warp-striped load) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadKeys( + bit_ordered_type (&keys)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + bit_ordered_type oob_item, + ::cuda::std::true_type is_full_tile, + ::cuda::std::true_type warp_striped) + { + LoadDirectWarpStriped(threadIdx.x, d_keys_in + block_offset, keys); + } + + /** + * Load a tile of keys (specialized for partial tile, warp-striped load) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadKeys( + bit_ordered_type (&keys)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + bit_ordered_type oob_item, + ::cuda::std::false_type is_full_tile, + ::cuda::std::true_type warp_striped) + { + // Register pressure work-around: moving valid_items through shfl prevents compiler + // from reusing guards/addressing from prior guarded loads + valid_items = ::cuda::device::warp_shuffle_idx(valid_items, 0); + + LoadDirectWarpStriped(threadIdx.x, d_keys_in + block_offset, keys, valid_items, oob_item); + } + + /** + * Load a tile of values (specialized for full tile, block load) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadValues( + ValueT (&values)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + ::cuda::std::true_type is_full_tile, + ::cuda::std::false_type warp_striped) + { + BlockLoadValuesT(temp_storage.load_values).Load(d_values_in + block_offset, values); + + __syncthreads(); + } + + /** + * Load a tile of values (specialized for partial tile, block load) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadValues( + ValueT (&values)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + ::cuda::std::false_type is_full_tile, + ::cuda::std::false_type warp_striped) + { + // Register pressure work-around: moving valid_items through shfl prevents compiler + // from reusing guards/addressing from prior guarded loads + valid_items = ::cuda::device::warp_shuffle_idx(valid_items, 0); + + BlockLoadValuesT(temp_storage.load_values).Load(d_values_in + block_offset, values, valid_items); + + __syncthreads(); + } + + /** + * Load a tile of items (specialized for full tile, warp-striped load) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadValues( + ValueT (&values)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + ::cuda::std::true_type is_full_tile, + ::cuda::std::true_type warp_striped) + { + LoadDirectWarpStriped(threadIdx.x, d_values_in + block_offset, values); + } + + /** + * Load a tile of items (specialized for partial tile, warp-striped load) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadValues( + ValueT (&values)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + ::cuda::std::false_type is_full_tile, + ::cuda::std::true_type warp_striped) + { + // Register pressure work-around: moving valid_items through shfl prevents compiler + // from reusing guards/addressing from prior guarded loads + valid_items = ::cuda::device::warp_shuffle_idx(valid_items, 0); + + LoadDirectWarpStriped(threadIdx.x, d_values_in + block_offset, values, valid_items); + } + + /** + * Truck along associated values + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void GatherScatterValues( + OffsetT (&relative_bin_offsets)[ITEMS_PER_THREAD], + int (&ranks)[ITEMS_PER_THREAD], + OffsetT block_offset, + OffsetT valid_items, + ::cuda::std::false_type /*is_keys_only*/) + { + ValueT values[ITEMS_PER_THREAD]; + + __syncthreads(); + + LoadValues(values, block_offset, valid_items, bool_constant_v, bool_constant_v); + + ScatterValues(values, relative_bin_offsets, ranks, valid_items); + } + + /** + * Truck along associated values (specialized for key-only sorting) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void GatherScatterValues( + OffsetT (& /*relative_bin_offsets*/)[ITEMS_PER_THREAD], + int (& /*ranks*/)[ITEMS_PER_THREAD], + OffsetT /*block_offset*/, + OffsetT /*valid_items*/, + ::cuda::std::true_type /*is_keys_only*/) + {} + + /** + * Process tile + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ProcessTile(OffsetT block_offset, OffsetT valid_items = TILE_ITEMS) + { + bit_ordered_type keys[ITEMS_PER_THREAD]; + int ranks[ITEMS_PER_THREAD]; + OffsetT relative_bin_offsets[ITEMS_PER_THREAD]; + + // Assign default (min/max) value to all keys + bit_ordered_type default_key = + IS_DESCENDING ? traits::min_raw_binary_key(decomposer) : traits::max_raw_binary_key(decomposer); + + // Load tile of keys + LoadKeys( + keys, block_offset, valid_items, default_key, bool_constant_v, bool_constant_v); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int KEY = 0; KEY < ITEMS_PER_THREAD; KEY++) + { + keys[KEY] = bit_ordered_conversion::to_bit_ordered(decomposer, keys[KEY]); + } + + // Rank the twiddled keys + int exclusive_digit_prefix[BINS_TRACKED_PER_THREAD]; + BlockRadixRankT(temp_storage.radix_rank).RankKeys(keys, ranks, digit_extractor(), exclusive_digit_prefix); + + __syncthreads(); + + // Share exclusive digit prefix + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + // Store exclusive prefix + temp_storage.exclusive_digit_prefix[bin_idx] = exclusive_digit_prefix[track]; + } + } + + __syncthreads(); + + // Get inclusive digit prefix + int inclusive_digit_prefix[BINS_TRACKED_PER_THREAD]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + if (IS_DESCENDING) + { + // Get inclusive digit prefix from exclusive prefix (higher bins come first) + inclusive_digit_prefix[track] = + (bin_idx == 0) ? (BLOCK_THREADS * ITEMS_PER_THREAD) : temp_storage.exclusive_digit_prefix[bin_idx - 1]; + } + else + { + // Get inclusive digit prefix from exclusive prefix (lower bins come first) + inclusive_digit_prefix[track] = + (bin_idx == RADIX_DIGITS - 1) + ? (BLOCK_THREADS * ITEMS_PER_THREAD) + : temp_storage.exclusive_digit_prefix[bin_idx + 1]; + } + } + } + + __syncthreads(); + + // Update global scatter base offsets for each digit + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + bin_offset[track] -= exclusive_digit_prefix[track]; + temp_storage.keys_and_offsets.relative_bin_offsets[bin_idx] = bin_offset[track]; + bin_offset[track] += inclusive_digit_prefix[track]; + } + } + + __syncthreads(); + + // Scatter keys + ScatterKeys(keys, relative_bin_offsets, ranks, valid_items); + + // Gather/scatter values + GatherScatterValues(relative_bin_offsets, ranks, block_offset, valid_items, bool_constant_v); + } + + //--------------------------------------------------------------------- + // Copy shortcut + //--------------------------------------------------------------------- + + /** + * Copy tiles within the range of input + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Copy(InputIteratorT d_in, T* d_out, OffsetT block_offset, OffsetT block_end) + { + // Simply copy the input + while (block_end - block_offset >= TILE_ITEMS) + { + T items[ITEMS_PER_THREAD]; + + LoadDirectStriped(threadIdx.x, d_in + block_offset, items); + __syncthreads(); + StoreDirectStriped(threadIdx.x, d_out + block_offset, items); + + block_offset += TILE_ITEMS; + } + + // Clean up last partial tile with guarded-I/O + if (block_offset < block_end) + { + OffsetT valid_items = block_end - block_offset; + + T items[ITEMS_PER_THREAD]; + + LoadDirectStriped(threadIdx.x, d_in + block_offset, items, valid_items); + __syncthreads(); + StoreDirectStriped(threadIdx.x, d_out + block_offset, items, valid_items); + } + } + + /** + * Copy tiles within the range of input (specialized for NullType) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + Copy(InputIteratorT /*d_in*/, NullType* /*d_out*/, OffsetT /*block_offset*/, OffsetT /*block_end*/) + {} + + //--------------------------------------------------------------------- + // Interface + //--------------------------------------------------------------------- + + /** + * Constructor + */ + _CCCL_DEVICE _CCCL_FORCEINLINE AgentRadixSortDownsweep( + TempStorage& temp_storage, + OffsetT (&bin_offset)[BINS_TRACKED_PER_THREAD], + OffsetT num_items, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + int current_bit, + int num_bits, + DecomposerT decomposer = {}) + : temp_storage(temp_storage.Alias()) + , d_keys_in(reinterpret_cast(d_keys_in)) + , d_values_in(d_values_in) + , d_keys_out(reinterpret_cast(d_keys_out)) + , d_values_out(d_values_out) + , current_bit(current_bit) + , num_bits(num_bits) + , short_circuit(1) + , decomposer(decomposer) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + this->bin_offset[track] = bin_offset[track]; + + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + // Short circuit if the histogram has only bin counts of only zeros or problem-size + short_circuit = short_circuit && ((bin_offset[track] == 0) || (bin_offset[track] == num_items)); + } + } + + short_circuit = __syncthreads_and(short_circuit); + } + + /** + * Constructor + */ + _CCCL_DEVICE _CCCL_FORCEINLINE AgentRadixSortDownsweep( + TempStorage& temp_storage, + OffsetT num_items, + OffsetT* d_spine, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + int current_bit, + int num_bits, + DecomposerT decomposer = {}) + : temp_storage(temp_storage.Alias()) + , d_keys_in(reinterpret_cast(d_keys_in)) + , d_values_in(d_values_in) + , d_keys_out(reinterpret_cast(d_keys_out)) + , d_values_out(d_values_out) + , current_bit(current_bit) + , num_bits(num_bits) + , short_circuit(1) + , decomposer(decomposer) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + + // Load digit bin offsets (each of the first RADIX_DIGITS threads will load an offset for that digit) + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + if (IS_DESCENDING) + { + bin_idx = RADIX_DIGITS - bin_idx - 1; + } + + // Short circuit if the first block's histogram has only bin counts of only zeros or problem-size + OffsetT first_block_bin_offset = d_spine[gridDim.x * bin_idx]; + short_circuit = short_circuit && ((first_block_bin_offset == 0) || (first_block_bin_offset == num_items)); + + // Load my block's bin offset for my bin + bin_offset[track] = d_spine[(gridDim.x * bin_idx) + blockIdx.x]; + } + } + + short_circuit = __syncthreads_and(short_circuit); + } + + /** + * Distribute keys from a segment of input tiles. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ProcessRegion(OffsetT block_offset, OffsetT block_end) + { + if (short_circuit) + { + // Copy keys + Copy(d_keys_in, d_keys_out, block_offset, block_end); + + // Copy values + Copy(d_values_in, d_values_out, block_offset, block_end); + } + else + { + // Process full tiles of tile_items + _CCCL_PRAGMA_NOUNROLL() + while (block_end - block_offset >= TILE_ITEMS) + { + ProcessTile(block_offset); + block_offset += TILE_ITEMS; + + __syncthreads(); + } + + // Clean up last partial tile with guarded-I/O + if (block_offset < block_end) + { + ProcessTile(block_offset, block_end - block_offset); + } + } + } +}; +} // namespace detail::radix_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_histogram.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_histogram.cuh new file mode 100644 index 00000000..b660e3c7 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_histogram.cuh @@ -0,0 +1,288 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2020, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * agent_radix_sort_histogram.cuh implements a stateful abstraction of CUDA + * thread blocks for participating in the device histogram kernel used for + * one-sweep radix sorting. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +//! @param ComputeT If void, use NOMINAL_4B_NUM_PARTS directly for NUM_PARTS. Otherwise, perform scaling. +template +struct agent_radix_sort_histogram_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int ITEMS_PER_THREAD = ItemsPerThread; + + // need to discard sizeof(ComputeType) in case it's void + template + _CCCL_HOST_DEVICE_API static constexpr int num_parts_helper() + { + if constexpr (::cuda::std::is_void_v) + { + return NOMINAL_4B_NUM_PARTS; + } + else + { + return ::cuda::std::max(1, NOMINAL_4B_NUM_PARTS * 4 / ::cuda::std::max(int{sizeof(ComputeType)}, 4)); + } + } + + /** NUM_PARTS is the number of private histograms (parts) each histogram is split + * into. Each warp lane is assigned to a specific part based on the lane + * ID. However, lanes with the same ID in different warp use the same private + * histogram. This arrangement helps reduce the degree of conflicts in atomic + * operations. */ + static constexpr int NUM_PARTS = num_parts_helper(); + + static constexpr int RADIX_BITS = RadixBits; +}; + +template +struct agent_radix_sort_exclusive_sum_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int RADIX_BITS = RadixBits; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template +using AgentRadixSortHistogramPolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceRadixSort") = + detail::agent_radix_sort_histogram_policy; + +//! Deprecated [Since 3.5] +template +using AgentRadixSortExclusiveSumPolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceRadixSort") = + detail::agent_radix_sort_exclusive_sum_policy; + +namespace detail::radix_sort +{ +template +struct AgentRadixSortHistogram +{ + // constants + static constexpr int ITEMS_PER_THREAD = AgentRadixSortHistogramPolicy::ITEMS_PER_THREAD; + static constexpr int BLOCK_THREADS = AgentRadixSortHistogramPolicy::BLOCK_THREADS; + static constexpr int TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD; + static constexpr int RADIX_BITS = AgentRadixSortHistogramPolicy::RADIX_BITS; + static constexpr int RADIX_DIGITS = 1 << RADIX_BITS; + static constexpr int MAX_NUM_PASSES = (sizeof(KeyT) * 8 + RADIX_BITS - 1) / RADIX_BITS; + static constexpr int NUM_PARTS = AgentRadixSortHistogramPolicy::NUM_PARTS; + + using traits = radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion = typename traits::bit_ordered_conversion_policy; + + using Twiddle = RadixSortTwiddle; + using ShmemCounterT = uint32_t; + using ShmemAtomicCounterT = ShmemCounterT; + + using fundamental_digit_extractor_t = ShiftDigitExtractor; + using digit_extractor_t = typename traits::template digit_extractor_t; + + struct _TempStorage + { + ShmemAtomicCounterT bins[MAX_NUM_PASSES][RADIX_DIGITS][NUM_PARTS]; + }; + + struct TempStorage : Uninitialized<_TempStorage> + {}; + + // thread fields + // shared memory storage + _TempStorage& s; + + // bins for the histogram + OffsetT* d_bins_out; + + // data to compute the histogram + const bit_ordered_type* d_keys_in; + + // number of data items + OffsetT num_items; + + // begin and end bits for sorting + int begin_bit, end_bit; + + // number of sorting passes + int num_passes; // NOLINT(modernize-use-default-member-init) + + DecomposerT decomposer; + + _CCCL_DEVICE _CCCL_FORCEINLINE AgentRadixSortHistogram( + TempStorage& temp_storage, + OffsetT* d_bins_out, + const KeyT* d_keys_in, + OffsetT num_items, + int begin_bit, + int end_bit, + DecomposerT decomposer = {}) + : s(temp_storage.Alias()) + , d_bins_out(d_bins_out) + , d_keys_in(reinterpret_cast(d_keys_in)) + , num_items(num_items) + , begin_bit(begin_bit) + , end_bit(end_bit) + , num_passes((end_bit - begin_bit + RADIX_BITS - 1) / RADIX_BITS) + , decomposer(decomposer) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE void Init() + { + // Initialize bins to 0. + _CCCL_PRAGMA_UNROLL_FULL() + for (int bin = static_cast(threadIdx.x); bin < RADIX_DIGITS; bin += BLOCK_THREADS) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int pass = 0; pass < num_passes; ++pass) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int part = 0; part < NUM_PARTS; ++part) + { + s.bins[pass][bin][part] = 0; + } + } + } + __syncthreads(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadTileKeys(OffsetT tile_offset, bit_ordered_type (&keys)[ITEMS_PER_THREAD]) + { + // tile_offset < num_items always, hence the line below works + bool full_tile = num_items - tile_offset >= TILE_ITEMS; + if (full_tile) + { + LoadDirectStriped(threadIdx.x, d_keys_in + tile_offset, keys); + } + else + { + LoadDirectStriped( + threadIdx.x, d_keys_in + tile_offset, keys, num_items - tile_offset, Twiddle::DefaultKey(decomposer)); + } + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + keys[u] = Twiddle::In(keys[u], decomposer); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + AccumulateSharedHistograms(OffsetT tile_offset, bit_ordered_type (&keys)[ITEMS_PER_THREAD]) + { + int part = ::cuda::ptx::get_sreg_laneid() % NUM_PARTS; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int current_bit = begin_bit, pass = 0; current_bit < end_bit; current_bit += RADIX_BITS, ++pass) + { + const int num_bits = ::cuda::std::min(+RADIX_BITS, end_bit - current_bit); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + uint32_t bin = digit_extractor(current_bit, num_bits).Digit(keys[u]); + // Using cuda::atomic<> results in lower performance on GP100, + // so atomicAdd() is used instead. + atomicAdd(&s.bins[pass][bin][part], 1); + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void AccumulateGlobalHistograms() + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int bin = static_cast(threadIdx.x); bin < RADIX_DIGITS; bin += BLOCK_THREADS) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int pass = 0; pass < num_passes; ++pass) + { + OffsetT count = cub::ThreadReduce(s.bins[pass][bin], ::cuda::std::plus<>{}); + if (count > 0) + { + // Using cuda::atomic<> here would also require using it in + // other kernels. However, other kernels of onesweep sorting + // (ExclusiveSum, Onesweep) don't need atomic + // access. Therefore, atomicAdd() is used, until + // cuda::atomic_ref<> becomes available. + atomicAdd(&d_bins_out[pass * RADIX_DIGITS + bin], count); + } + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void Process() + { + _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); + // Within a portion, avoid overflowing (u)int32 counters. + // Between portions, accumulate results in global memory. + constexpr OffsetT MAX_PORTION_SIZE = 1 << 30; + OffsetT num_portions = ::cuda::ceil_div(num_items, MAX_PORTION_SIZE); + for (OffsetT portion = 0; portion < num_portions; ++portion) + { + // Reset the counters. + Init(); + + // Process the tiles. + OffsetT portion_offset = portion * MAX_PORTION_SIZE; + OffsetT portion_size = ::cuda::std::min(MAX_PORTION_SIZE, num_items - portion_offset); + for (OffsetT offset = static_cast(blockIdx.x) * TILE_ITEMS; offset < portion_size; + offset += OffsetT{TILE_ITEMS} * gridDim.x) + { + OffsetT tile_offset = portion_offset + offset; + bit_ordered_type keys[ITEMS_PER_THREAD]; + LoadTileKeys(tile_offset, keys); + AccumulateSharedHistograms(tile_offset, keys); + } + __syncthreads(); + + // Accumulate the result in global memory. + // Wait for global histogram init + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + AccumulateGlobalHistograms(); + __syncthreads(); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE digit_extractor_t digit_extractor(int current_bit, int num_bits) + { + return traits::template digit_extractor(current_bit, num_bits, decomposer); + } +}; +} // namespace detail::radix_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_onesweep.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_onesweep.cuh new file mode 100644 index 00000000..23bb7601 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_onesweep.cuh @@ -0,0 +1,742 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2020, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * agent_radix_sort_onesweep.cuh implements a stateful abstraction of CUDA + * thread blocks for participating in the device one-sweep radix sort kernel. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/** \brief cub::RadixSortStoreAlgorithm enumerates different algorithms to write + * partitioned elements (keys, values) stored in shared memory into global + * memory. Currently applies only to writing 4B keys in full tiles; in all other cases, + * RADIX_SORT_STORE_DIRECT is used. + */ +enum RadixSortStoreAlgorithm +{ + /** \brief Elements are statically distributed among block threads, which write them + * into the appropriate partition in global memory. This results in fewer instructions + * and more writes in flight at a given moment, but may generate more transactions. */ + RADIX_SORT_STORE_DIRECT, + /** \brief Elements are distributed among warps in a block distribution. Each warp + * goes through its elements and tries to write them while minimizing the number of + * memory transactions. This results in fewer memory transactions, but more + * instructions and less writes in flight at a given moment. */ + RADIX_SORT_STORE_ALIGNED +}; + +#if _CCCL_HOSTED() +namespace detail +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(RadixSortStoreAlgorithm algo) noexcept +{ + switch (algo) + { + case RADIX_SORT_STORE_DIRECT: + return "RADIX_SORT_STORE_DIRECT"; + case RADIX_SORT_STORE_ALIGNED: + return "RADIX_SORT_STORE_ALIGNED"; + } + return ""; +} +} // namespace detail +#endif // _CCCL_HOSTED() + +#if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED) +inline ::std::ostream& operator<<(::std::ostream& os, RadixSortStoreAlgorithm algo) +{ + return os << CUB_NS_QUALIFIER::detail::to_string(algo); +} +#endif // _CCCL_HOSTED() && !_CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END + +#if __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) +template <::cuda::std::same_as CharT> +struct std::formatter : formatter +{ + template + auto format(const CUB_NS_QUALIFIER::RadixSortStoreAlgorithm& algo, FmtCtx& ctx) const + { + return formatter::format(CUB_NS_QUALIFIER::detail::to_string(algo), ctx); + } +}; +#endif // __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +template > +struct agent_radix_sort_onesweep_policy : ScalingType +{ + static constexpr int RANK_NUM_PARTS = RankNumParts; + static constexpr int RADIX_BITS = RadixBits; + static constexpr RadixRankAlgorithm RANK_ALGORITHM = RankAlgorithm; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + static constexpr RadixSortStoreAlgorithm STORE_ALGORITHM = StoreAlgorithm; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentRadixSortOnesweepPolicy + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceRadixSort") = detail::agent_radix_sort_onesweep_policy< + NominalThreadsPerBlock4B, + NominalItemsPerThread4B, + ComputeT, + RankNumParts, + RankAlgorithm, + ScanAlgorithm, + StoreAlgorithm, + RadixBits, + ScalingType>; + +namespace detail::radix_sort +{ +template +struct AgentRadixSortOnesweep +{ + // constants + static constexpr int ITEMS_PER_THREAD = AgentRadixSortOnesweepPolicy::ITEMS_PER_THREAD; + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + static constexpr int BLOCK_THREADS = AgentRadixSortOnesweepPolicy::BLOCK_THREADS; + static constexpr int RANK_NUM_PARTS = AgentRadixSortOnesweepPolicy::RANK_NUM_PARTS; + static constexpr int TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD; + static constexpr int RADIX_BITS = AgentRadixSortOnesweepPolicy::RADIX_BITS; + static constexpr int RADIX_DIGITS = 1 << RADIX_BITS; + static constexpr int BINS_PER_THREAD = (RADIX_DIGITS + BLOCK_THREADS - 1) / BLOCK_THREADS; + static constexpr bool FULL_BINS = BINS_PER_THREAD * BLOCK_THREADS == RADIX_DIGITS; + static constexpr int WARP_THREADS = warp_threads; + static constexpr int BLOCK_WARPS = BLOCK_THREADS / WARP_THREADS; + static constexpr int WARP_MASK = ~0; + static constexpr int LOOKBACK_PARTIAL_MASK = 1 << (PortionOffsetT(sizeof(PortionOffsetT)) * 8 - 2); + static constexpr int LOOKBACK_GLOBAL_MASK = 1 << (PortionOffsetT(sizeof(PortionOffsetT)) * 8 - 1); + static constexpr int LOOKBACK_KIND_MASK = LOOKBACK_PARTIAL_MASK | LOOKBACK_GLOBAL_MASK; + static constexpr int LOOKBACK_VALUE_MASK = ~LOOKBACK_KIND_MASK; + + using traits = radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion = typename traits::bit_ordered_conversion_policy; + + using fundamental_digit_extractor_t = ShiftDigitExtractor; + using digit_extractor_t = typename traits::template digit_extractor_t; + + using AtomicOffsetT = PortionOffsetT; + + static constexpr RadixRankAlgorithm RANK_ALGORITHM = AgentRadixSortOnesweepPolicy::RANK_ALGORITHM; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = AgentRadixSortOnesweepPolicy::SCAN_ALGORITHM; + static constexpr RadixSortStoreAlgorithm STORE_ALGORITHM = + sizeof(bit_ordered_type) == sizeof(uint32_t) + ? AgentRadixSortOnesweepPolicy::STORE_ALGORITHM + : RADIX_SORT_STORE_DIRECT; + + using Twiddle = RadixSortTwiddle; + + static_assert(RANK_ALGORITHM == RADIX_RANK_MATCH || RANK_ALGORITHM == RADIX_RANK_MATCH_EARLY_COUNTS_ANY + || RANK_ALGORITHM == RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR, + "for onesweep agent, the ranking algorithm must warp-strided key arrangement"); + + using BlockRadixRankT = ::cuda::std::_If< + RANK_ALGORITHM == RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR, + BlockRadixRankMatchEarlyCounts, + ::cuda::std::_If< + RANK_ALGORITHM == RADIX_RANK_MATCH, + BlockRadixRankMatch, + BlockRadixRankMatchEarlyCounts>>; + + // temporary storage + struct TempStorage_ + { + union + { + bit_ordered_type keys_out[TILE_ITEMS]; + ValueT values_out[TILE_ITEMS]; + typename BlockRadixRankT::TempStorage rank_temp_storage; + }; + union + { + OffsetT global_offsets[RADIX_DIGITS]; + PortionOffsetT block_idx; + }; + }; + + using TempStorage = Uninitialized; + + // thread variables + TempStorage_& s; + + // kernel parameters + AtomicOffsetT* d_lookback; + AtomicOffsetT* d_ctrs; + OffsetT* d_bins_out; + const OffsetT* d_bins_in; + bit_ordered_type* d_keys_out; + const bit_ordered_type* d_keys_in; + ValueT* d_values_out; + const ValueT* d_values_in; + PortionOffsetT num_items; + int current_bit; + int num_bits; + + // other thread variables + int warp; + int lane; + DecomposerT decomposer; + PortionOffsetT block_idx; + bool full_block; + + _CCCL_DEVICE _CCCL_FORCEINLINE digit_extractor_t digit_extractor() + { + return traits::template digit_extractor(current_bit, num_bits, decomposer); + } + + // helper methods + _CCCL_DEVICE _CCCL_FORCEINLINE uint32_t Digit(bit_ordered_type key) + { + return digit_extractor().Digit(key); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE int ThreadBin(int u) + { + return threadIdx.x * BINS_PER_THREAD + u; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void LookbackPartial(int (&bins)[BINS_PER_THREAD]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + int bin = ThreadBin(u); + if (FULL_BINS || bin < RADIX_DIGITS) + { + // write the local sum into the bin + AtomicOffsetT& loc = d_lookback[block_idx * RADIX_DIGITS + bin]; + PortionOffsetT value = bins[u] | LOOKBACK_PARTIAL_MASK; + ThreadStore(&loc, value); + } + } + } + + struct CountsCallback + { + using AgentT = + AgentRadixSortOnesweep; + AgentT& agent; + int (&bins)[BINS_PER_THREAD]; + bit_ordered_type (&keys)[ITEMS_PER_THREAD]; + static constexpr bool EMPTY = false; + _CCCL_DEVICE _CCCL_FORCEINLINE + CountsCallback(AgentT& agent, int (&bins)[BINS_PER_THREAD], bit_ordered_type (&keys)[ITEMS_PER_THREAD]) + : agent(agent) + , bins(bins) + , keys(keys) + {} + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(int (&other_bins)[BINS_PER_THREAD]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + bins[u] = other_bins[u]; + } + + // Wait for lookback init + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + agent.LookbackPartial(bins); + + agent.TryShortCircuit(keys, bins); + } + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE void LookbackGlobal(int (&bins)[BINS_PER_THREAD]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + int bin = ThreadBin(u); + if (FULL_BINS || bin < RADIX_DIGITS) + { + PortionOffsetT inc_sum = bins[u]; + int want_mask = ~0; + // backtrack as long as necessary + for (PortionOffsetT block_jdx = block_idx - 1; block_jdx >= 0; --block_jdx) + { + // wait for some value to appear + PortionOffsetT value_j = 0; + AtomicOffsetT& loc_j = d_lookback[block_jdx * RADIX_DIGITS + bin]; + do + { + __threadfence_block(); // prevent hoisting loads from loop + value_j = ThreadLoad(&loc_j); + } while (value_j == 0); + + inc_sum += value_j & LOOKBACK_VALUE_MASK; + want_mask = __ballot_sync(want_mask, (value_j & LOOKBACK_GLOBAL_MASK) == 0); + if (value_j & LOOKBACK_GLOBAL_MASK) + { + break; + } + } + AtomicOffsetT& loc_i = d_lookback[block_idx * RADIX_DIGITS + bin]; + PortionOffsetT value_i = inc_sum | LOOKBACK_GLOBAL_MASK; + ThreadStore(&loc_i, value_i); + s.global_offsets[bin] += inc_sum - bins[u]; + } + } + _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadKeys(OffsetT tile_offset, bit_ordered_type (&keys)[ITEMS_PER_THREAD]) + { + if (full_block) + { + LoadDirectWarpStriped(threadIdx.x, d_keys_in + tile_offset, keys); + } + else + { + LoadDirectWarpStriped( + threadIdx.x, d_keys_in + tile_offset, keys, num_items - tile_offset, Twiddle::DefaultKey(decomposer)); + } + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + keys[u] = Twiddle::In(keys[u], decomposer); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadValues(OffsetT tile_offset, ValueT (&values)[ITEMS_PER_THREAD]) + { + if (full_block) + { + LoadDirectWarpStriped(threadIdx.x, d_values_in + tile_offset, values); + } + else + { + int tile_items = num_items - tile_offset; + LoadDirectWarpStriped(threadIdx.x, d_values_in + tile_offset, values, tile_items); + } + } + + /** Checks whether "short-circuiting" is possible. Short-circuiting happens + * if all TILE_ITEMS keys fall into the same bin, i.e. have the same digit + * value (note that it only happens for full tiles). If short-circuiting is + * performed, the part of the ranking algorithm after the CountsCallback, as + * well as the rest of the sorting (e.g. scattering keys and values to + * shared and global memory) are skipped; updates related to decoupled + * look-back are still performed. Instead, the keys assigned to the current + * thread block are written cooperatively into a contiguous location in + * d_keys_out corresponding to their digit. The values (if also sorting + * values) assigned to the current thread block are similarly copied from + * d_values_in to d_values_out. */ + _CCCL_DEVICE _CCCL_FORCEINLINE void + TryShortCircuit(bit_ordered_type (&keys)[ITEMS_PER_THREAD], int (&bins)[BINS_PER_THREAD]) + { + // check if any bin can be short-circuited + bool short_circuit = false; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + if (FULL_BINS || ThreadBin(u) < RADIX_DIGITS) + { + short_circuit = short_circuit || bins[u] == TILE_ITEMS; + } + } + short_circuit = __syncthreads_or(short_circuit); + if (!short_circuit) + { + return; + } + + ShortCircuitCopy(keys, bins); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + ShortCircuitCopy(bit_ordered_type (&keys)[ITEMS_PER_THREAD], int (&bins)[BINS_PER_THREAD]) + { + // short-circuit handling; note that global look-back is still required + + // compute offsets + uint32_t common_bin = Digit(keys[0]); + int offsets[BINS_PER_THREAD]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + int bin = ThreadBin(u); + offsets[u] = bin > common_bin ? TILE_ITEMS : 0; + } + + // global lookback + LoadBinsToOffsetsGlobal(offsets); + LookbackGlobal(bins); + UpdateBinsGlobal(bins, offsets); + __syncthreads(); + + // scatter the keys + OffsetT global_offset = s.global_offsets[common_bin]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + keys[u] = Twiddle::Out(keys[u], decomposer); + } + if (full_block) + { + StoreDirectWarpStriped(threadIdx.x, d_keys_out + global_offset, keys); + } + else + { + int tile_items = num_items - block_idx * TILE_ITEMS; + StoreDirectWarpStriped(threadIdx.x, d_keys_out + global_offset, keys, tile_items); + } + + if (!KEYS_ONLY) + { + // gather and scatter the values + ValueT values[ITEMS_PER_THREAD]; + LoadValues(block_idx * TILE_ITEMS, values); // NOLINT(bugprone-misplaced-widening-cast) + if (full_block) + { + StoreDirectWarpStriped(threadIdx.x, d_values_out + global_offset, values); + } + else + { + int tile_items = num_items - block_idx * TILE_ITEMS; + StoreDirectWarpStriped(threadIdx.x, d_values_out + global_offset, values, tile_items); + } + } + + // exit early + ThreadExit(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + ScatterKeysShared(bit_ordered_type (&keys)[ITEMS_PER_THREAD], int (&ranks)[ITEMS_PER_THREAD]) + { + // write to shared memory + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + s.keys_out[ranks[u]] = keys[u]; + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + ScatterValuesShared(ValueT (&values)[ITEMS_PER_THREAD], int (&ranks)[ITEMS_PER_THREAD]) + { + // write to shared memory + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + s.values_out[ranks[u]] = values[u]; + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void LoadBinsToOffsetsGlobal(int (&offsets)[BINS_PER_THREAD]) + { + // global offset - global part + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + int bin = ThreadBin(u); + if (FULL_BINS || bin < RADIX_DIGITS) + { + s.global_offsets[bin] = d_bins_in[bin] - offsets[u]; + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void UpdateBinsGlobal(int (&bins)[BINS_PER_THREAD], int (&offsets)[BINS_PER_THREAD]) + { + bool last_block = (block_idx + 1) * TILE_ITEMS >= num_items; + if (d_bins_out != nullptr && last_block) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + int bin = ThreadBin(u); + if (FULL_BINS || bin < RADIX_DIGITS) + { + d_bins_out[bin] = s.global_offsets[bin] + offsets[u] + bins[u]; + } + } + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterKeysGlobalDirect() + { + int tile_items = FULL_TILE ? TILE_ITEMS : num_items - block_idx * TILE_ITEMS; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + int idx = threadIdx.x + u * BLOCK_THREADS; + bit_ordered_type key = s.keys_out[idx]; + OffsetT global_idx = idx + s.global_offsets[Digit(key)]; + if (FULL_TILE || idx < tile_items) + { + d_keys_out[global_idx] = Twiddle::Out(key, decomposer); + } + __syncwarp(WARP_MASK); + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterValuesGlobalDirect(int (&digits)[ITEMS_PER_THREAD]) + { + int tile_items = FULL_TILE ? TILE_ITEMS : num_items - block_idx * TILE_ITEMS; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + int idx = threadIdx.x + u * BLOCK_THREADS; + ValueT value = s.values_out[idx]; + OffsetT global_idx = idx + s.global_offsets[digits[u]]; + if (FULL_TILE || idx < tile_items) + { + d_values_out[global_idx] = value; + } + __syncwarp(WARP_MASK); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterKeysGlobalAligned() + { + // this only works with full tiles + constexpr int ITEMS_PER_WARP = TILE_ITEMS / BLOCK_WARPS; + constexpr int ALIGN = 8; + constexpr auto CACHE_MODIFIER = STORE_CG; + + int warp_start = warp * ITEMS_PER_WARP; + int warp_end = (warp + 1) * ITEMS_PER_WARP; + int warp_offset = warp_start; + while (warp_offset < warp_end - WARP_THREADS) + { + int idx = warp_offset + lane; + bit_ordered_type key = s.keys_out[idx]; + bit_ordered_type key_out = Twiddle::Out(key, decomposer); + OffsetT global_idx = idx + s.global_offsets[Digit(key)]; + int last_lane = WARP_THREADS - 1; + int num_writes = WARP_THREADS; + if (lane == last_lane) + { + num_writes -= int(global_idx + 1) % ALIGN; + } + num_writes = __shfl_sync(WARP_MASK, num_writes, last_lane); + if (lane < num_writes) + { + ThreadStore(&d_keys_out[global_idx], key_out); + } + warp_offset += num_writes; + } + { + int num_writes = warp_end - warp_offset; + if (lane < num_writes) + { + int idx = warp_offset + lane; + bit_ordered_type key = s.keys_out[idx]; + OffsetT global_idx = idx + s.global_offsets[Digit(key)]; + ThreadStore(&d_keys_out[global_idx], Twiddle::Out(key, decomposer)); + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterKeysGlobal() + { + // write block data to global memory + if (full_block) + { + if constexpr (STORE_ALGORITHM == RADIX_SORT_STORE_ALIGNED) + { + ScatterKeysGlobalAligned(); + } + else + { + ScatterKeysGlobalDirect(); + } + } + else + { + ScatterKeysGlobalDirect(); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterValuesGlobal(int (&digits)[ITEMS_PER_THREAD]) + { + // write block data to global memory + if (full_block) + { + ScatterValuesGlobalDirect(digits); + } + else + { + ScatterValuesGlobalDirect(digits); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ComputeKeyDigits(int (&digits)[ITEMS_PER_THREAD]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < ITEMS_PER_THREAD; ++u) + { + int idx = threadIdx.x + u * BLOCK_THREADS; + digits[u] = Digit(s.keys_out[idx]); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + GatherScatterValues(int (&ranks)[ITEMS_PER_THREAD], ::cuda::std::false_type keys_only) + { + // compute digits corresponding to the keys + int digits[ITEMS_PER_THREAD]; + ComputeKeyDigits(digits); + + // load values + ValueT values[ITEMS_PER_THREAD]; + LoadValues(block_idx * TILE_ITEMS, values); // NOLINT(bugprone-misplaced-widening-cast) + + // scatter values + __syncthreads(); + ScatterValuesShared(values, ranks); + + __syncthreads(); + ScatterValuesGlobal(digits); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + GatherScatterValues(int (&ranks)[ITEMS_PER_THREAD], ::cuda::std::true_type keys_only) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE void Process() + { + // load keys + // if warp1 < warp2, all elements of warp1 occur before those of warp2 + // in the source array + bit_ordered_type keys[ITEMS_PER_THREAD]; + LoadKeys(block_idx * TILE_ITEMS, keys); // NOLINT(bugprone-misplaced-widening-cast) + + // rank keys + int ranks[ITEMS_PER_THREAD]; + int exclusive_digit_prefix[BINS_PER_THREAD]; + int bins[BINS_PER_THREAD]; + BlockRadixRankT(s.rank_temp_storage) + .RankKeys(keys, ranks, digit_extractor(), exclusive_digit_prefix, CountsCallback(*this, bins, keys)); + + // scatter keys in shared memory + __syncthreads(); + ScatterKeysShared(keys, ranks); + + // compute global offsets + LoadBinsToOffsetsGlobal(exclusive_digit_prefix); + LookbackGlobal(bins); + UpdateBinsGlobal(bins, exclusive_digit_prefix); + + // scatter keys in global memory + __syncthreads(); + ScatterKeysGlobal(); + + // scatter values if necessary + GatherScatterValues(ranks, bool_constant_v); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE // + AgentRadixSortOnesweep( + TempStorage& temp_storage, + AtomicOffsetT* d_lookback, + AtomicOffsetT* d_ctrs, + OffsetT* d_bins_out, + const OffsetT* d_bins_in, + KeyT* d_keys_out, + const KeyT* d_keys_in, + ValueT* d_values_out, + const ValueT* d_values_in, + PortionOffsetT num_items, + int current_bit, + int num_bits, + DecomposerT decomposer = {}) + : s(temp_storage.Alias()) + , d_lookback(d_lookback) + , d_ctrs(d_ctrs) + , d_bins_out(d_bins_out) + , d_bins_in(d_bins_in) + , d_keys_out(reinterpret_cast(d_keys_out)) + , d_keys_in(reinterpret_cast(d_keys_in)) + , d_values_out(d_values_out) + , d_values_in(d_values_in) + , num_items(num_items) + , current_bit(current_bit) + , num_bits(num_bits) + , warp(static_cast(threadIdx.x / WARP_THREADS)) + , lane(static_cast(::cuda::ptx::get_sreg_laneid())) + , decomposer(decomposer) + { + // initialization + if (threadIdx.x == 0) + { + s.block_idx = atomicAdd(d_ctrs, 1); + } + __syncthreads(); + block_idx = s.block_idx; + full_block = (block_idx + 1) * TILE_ITEMS <= num_items; + } +}; +} // namespace detail::radix_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_upsweep.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_upsweep.cuh new file mode 100644 index 00000000..40fb9129 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_radix_sort_upsweep.cuh @@ -0,0 +1,517 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * AgentRadixSortUpsweep implements a stateful abstraction of CUDA thread blocks for participating in device-wide radix + * sort upsweep . + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +namespace detail +{ +/** + * @brief Parameterizable tuning policy type for AgentRadixSortUpsweep + * + * @tparam NominalThreadsPerBlock4B + * Threads per thread block + * + * @tparam NominalItemsPerThread4B + * Items per thread (per tile of input) + * + * @tparam ComputeT + * Dominant compute type + * + * @tparam LoadModifier + * Cache load modifier for reading keys + * + * @tparam RadixBits + * The number of radix bits, i.e., log2(bins) + */ +template > +struct agent_radix_sort_upsweep_policy : ScalingType +{ + /// The number of radix bits, i.e., log2(bins) + static constexpr int RADIX_BITS = RadixBits; + + /// Cache load modifier for reading keys + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentRadixSortUpsweepPolicy + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceRadixSort") = detail::agent_radix_sort_upsweep_policy< + NominalThreadsPerBlock4B, + NominalItemsPerThread4B, + ComputeT, + LoadModifier, + RadixBits, + ScalingType>; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::radix_sort +{ +/** + * @brief AgentRadixSortUpsweep implements a stateful abstraction of CUDA thread blocks for + * participating in device-wide radix sort upsweep . + * + * @tparam AgentRadixSortUpsweepPolicy + * Parameterized AgentRadixSortUpsweepPolicy tuning policy type + * + * @tparam KeyT + * KeyT type + * + * @tparam DecomposerT = identity_decomposer_t + * Signed integer type for global offsets + */ +template +struct AgentRadixSortUpsweep +{ + //--------------------------------------------------------------------- + // Type definitions and constants + //--------------------------------------------------------------------- + using traits = radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion = typename traits::bit_ordered_conversion_policy; + + // Integer type for digit counters (to be packed into words of PackedCounters) + using DigitCounter = unsigned char; + + // Integer type for packing DigitCounters into columns of shared memory banks + using PackedCounter = unsigned int; + + static constexpr CacheLoadModifier LOAD_MODIFIER = AgentRadixSortUpsweepPolicy::LOAD_MODIFIER; + + static constexpr int RADIX_BITS = AgentRadixSortUpsweepPolicy::RADIX_BITS; + static constexpr int BLOCK_THREADS = AgentRadixSortUpsweepPolicy::BLOCK_THREADS; + static constexpr int KEYS_PER_THREAD = AgentRadixSortUpsweepPolicy::ITEMS_PER_THREAD; + + static constexpr int RADIX_DIGITS = 1 << RADIX_BITS; + + static constexpr int LOG_WARP_THREADS = log2_warp_threads; + static constexpr int WARP_THREADS = 1 << LOG_WARP_THREADS; + static constexpr int WARPS = (BLOCK_THREADS + WARP_THREADS - 1) / WARP_THREADS; + + static constexpr int TILE_ITEMS = BLOCK_THREADS * KEYS_PER_THREAD; + + static constexpr int BYTES_PER_COUNTER = sizeof(DigitCounter); + static constexpr int OG_BYTES_PER_COUNTER = Log2::VALUE; + + static constexpr int PACKING_RATIO = sizeof(PackedCounter) / sizeof(DigitCounter); + static constexpr int LOG_PACKING_RATIO = Log2::VALUE; + + static constexpr int LOG_COUNTER_LANES = ::cuda::std::max(0, int(RADIX_BITS) - int(LOG_PACKING_RATIO)); + static constexpr int COUNTER_LANES = 1 << LOG_COUNTER_LANES; + + // To prevent counter overflow, we must periodically unpack and aggregate the + // digit counters back into registers. Each counter lane is assigned to a + // warp for aggregation. + + static constexpr int LANES_PER_WARP = ::cuda::std::max(1, (COUNTER_LANES + WARPS - 1) / WARPS); + + // Unroll tiles in batches without risk of counter overflow + static constexpr int UNROLL_COUNT = ::cuda::std::min(64, 255 / KEYS_PER_THREAD); + static constexpr int UNROLLED_ELEMENTS = UNROLL_COUNT * TILE_ITEMS; + + // Input iterator wrapper type (for applying cache modifier)s + using KeysItr = CacheModifiedInputIterator; + + // Digit extractor type + using fundamental_digit_extractor_t = BFEDigitExtractor; + using digit_extractor_t = typename traits::template digit_extractor_t; + + /** + * Shared memory storage layout + */ + union __align__(16) _TempStorage + { + DigitCounter thread_counters[COUNTER_LANES][BLOCK_THREADS][PACKING_RATIO]; + PackedCounter packed_thread_counters[COUNTER_LANES][BLOCK_THREADS]; + OffsetT block_counters[WARP_THREADS][RADIX_DIGITS]; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Thread fields (aggregate state bundle) + //--------------------------------------------------------------------- + + // Shared storage for this CTA + _TempStorage& temp_storage; + + // Thread-local counters for periodically aggregating composite-counter lanes + OffsetT local_counts[LANES_PER_WARP][PACKING_RATIO]; + + // Input and output device pointers + KeysItr d_keys_in; + + // Target bits + int current_bit; + int num_bits; + DecomposerT decomposer; + + //--------------------------------------------------------------------- + // Utility methods + //--------------------------------------------------------------------- + _CCCL_DEVICE _CCCL_FORCEINLINE digit_extractor_t digit_extractor() + { + return traits::template digit_extractor(current_bit, num_bits, decomposer); + } + + /** + * Decode a key and increment corresponding smem digit counter + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void Bucket(bit_ordered_type key) + { + // Perform transform op + bit_ordered_type converted_key = bit_ordered_conversion::to_bit_ordered(decomposer, key); + + // Extract current digit bits + uint32_t digit = digit_extractor().Digit(converted_key); + + // Get sub-counter offset + uint32_t sub_counter = digit & (PACKING_RATIO - 1); + + // Get row offset + uint32_t row_offset = digit >> LOG_PACKING_RATIO; + _CCCL_ASSERT(row_offset < COUNTER_LANES, ""); + + // Increment counter + temp_storage.thread_counters[row_offset][threadIdx.x][sub_counter]++; + } + + /** + * Reset composite counters + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ResetDigitCounters() + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int LANE = 0; LANE < COUNTER_LANES; LANE++) + { + temp_storage.packed_thread_counters[LANE][threadIdx.x] = 0; + } + } + + /** + * Reset the unpacked counters in each thread + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ResetUnpackedCounters() + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int LANE = 0; LANE < LANES_PER_WARP; LANE++) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int UNPACKED_COUNTER = 0; UNPACKED_COUNTER < PACKING_RATIO; UNPACKED_COUNTER++) + { + local_counts[LANE][UNPACKED_COUNTER] = 0; + } + } + } + + /** + * Extracts and aggregates the digit counters for each counter lane + * owned by this warp + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void UnpackDigitCounts() + { + unsigned int warp_id = threadIdx.x >> LOG_WARP_THREADS; + unsigned int warp_tid = ::cuda::ptx::get_sreg_laneid(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int LANE = 0; LANE < LANES_PER_WARP; LANE++) + { + const int counter_lane = (LANE * WARPS) + warp_id; + if (counter_lane < COUNTER_LANES) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int PACKED_COUNTER = 0; PACKED_COUNTER < BLOCK_THREADS; PACKED_COUNTER += WARP_THREADS) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int UNPACKED_COUNTER = 0; UNPACKED_COUNTER < PACKING_RATIO; UNPACKED_COUNTER++) + { + OffsetT counter = temp_storage.thread_counters[counter_lane][warp_tid + PACKED_COUNTER][UNPACKED_COUNTER]; + local_counts[LANE][UNPACKED_COUNTER] += counter; + } + } + } + } + } + + /** + * Processes a single, full tile + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ProcessFullTile(OffsetT block_offset) + { + // Tile of keys + bit_ordered_type keys[KEYS_PER_THREAD]; + + LoadDirectStriped(threadIdx.x, d_keys_in + block_offset, keys); + + // Prevent hoisting + __syncthreads(); + + // Bucket tile of keys + cuda::static_for([&](auto ic) { + Bucket(keys[ic]); + }); + } + + /** + * Processes a single load (may have some threads masked off) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ProcessPartialTile(OffsetT block_offset, const OffsetT& block_end) + { + // Process partial tile if necessary using single loads + for (OffsetT offset = threadIdx.x; offset < block_end - block_offset; offset += BLOCK_THREADS) + { + // Load and bucket key + bit_ordered_type key = d_keys_in[block_offset + offset]; + Bucket(key); + } + } + + //--------------------------------------------------------------------- + // Interface + //--------------------------------------------------------------------- + + /** + * Constructor + */ + _CCCL_DEVICE _CCCL_FORCEINLINE AgentRadixSortUpsweep( + TempStorage& temp_storage, const KeyT* d_keys_in, int current_bit, int num_bits, DecomposerT decomposer = {}) + : temp_storage(temp_storage.Alias()) + , d_keys_in(reinterpret_cast(d_keys_in)) + , current_bit(current_bit) + , num_bits(num_bits) + , decomposer(decomposer) + {} + + /** + * Compute radix digit histograms from a segment of input tiles. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ProcessRegion(OffsetT block_offset, const OffsetT& block_end) + { + // Reset digit counters in smem and unpacked counters in registers + ResetDigitCounters(); + ResetUnpackedCounters(); + + // Unroll batches of full tiles + while (block_end - block_offset >= UNROLLED_ELEMENTS) + { + for (int i = 0; i < UNROLL_COUNT; ++i) + { + ProcessFullTile(block_offset); + block_offset += TILE_ITEMS; + } + + __syncthreads(); + + // Aggregate back into local_count registers to prevent overflow + UnpackDigitCounts(); + + __syncthreads(); + + // Reset composite counters in lanes + ResetDigitCounters(); + } + + // Unroll single full tiles + while (block_end - block_offset >= TILE_ITEMS) + { + ProcessFullTile(block_offset); + block_offset += TILE_ITEMS; + } + + // Process partial tile if necessary + ProcessPartialTile(block_offset, block_end); + + __syncthreads(); + + // Aggregate back into local_count registers + UnpackDigitCounts(); + } + + /** + * Extract counts (saving them to the external array) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExtractCounts(OffsetT* counters, int bin_stride = 1, int bin_offset = 0) + { + unsigned int warp_id = threadIdx.x >> LOG_WARP_THREADS; + unsigned int warp_tid = ::cuda::ptx::get_sreg_laneid(); + + // Place unpacked digit counters in shared memory + _CCCL_PRAGMA_UNROLL_FULL() + for (int LANE = 0; LANE < LANES_PER_WARP; LANE++) + { + int counter_lane = (LANE * WARPS) + warp_id; + if (counter_lane < COUNTER_LANES) + { + int digit_row = counter_lane << LOG_PACKING_RATIO; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int UNPACKED_COUNTER = 0; UNPACKED_COUNTER < PACKING_RATIO; UNPACKED_COUNTER++) + { + int bin_idx = digit_row + UNPACKED_COUNTER; + + temp_storage.block_counters[warp_tid][bin_idx] = local_counts[LANE][UNPACKED_COUNTER]; + } + } + } + + __syncthreads(); + + // Rake-reduce bin_count reductions + + // Whole blocks + _CCCL_PRAGMA_UNROLL_FULL() + for (int BIN_BASE = RADIX_DIGITS % BLOCK_THREADS; (BIN_BASE + BLOCK_THREADS) <= RADIX_DIGITS; + BIN_BASE += BLOCK_THREADS) + { + int bin_idx = static_cast(BIN_BASE + threadIdx.x); + OffsetT bin_count = 0; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < WARP_THREADS; ++i) + { + bin_count += temp_storage.block_counters[i][bin_idx]; + } + + if (IS_DESCENDING) + { + bin_idx = RADIX_DIGITS - bin_idx - 1; + } + + counters[(bin_stride * bin_idx) + bin_offset] = bin_count; + } + + // Remainder + if ((RADIX_DIGITS % BLOCK_THREADS != 0) && (threadIdx.x < RADIX_DIGITS)) + { + int bin_idx = static_cast(threadIdx.x); + OffsetT bin_count = 0; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < WARP_THREADS; ++i) + { + bin_count += temp_storage.block_counters[i][bin_idx]; + } + + if (IS_DESCENDING) + { + bin_idx = RADIX_DIGITS - bin_idx - 1; + } + + counters[(bin_stride * bin_idx) + bin_offset] = bin_count; + } + } + + /** + * @brief Extract counts + * + * @param[out] bin_count + * The exclusive prefix sum for the digits + * [(threadIdx.x * BINS_TRACKED_PER_THREAD) ... (threadIdx.x * BINS_TRACKED_PER_THREAD) + BINS_TRACKED_PER_THREAD - + * 1] + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExtractCounts(OffsetT (&bin_count)[BINS_TRACKED_PER_THREAD]) + { + unsigned int warp_id = threadIdx.x >> LOG_WARP_THREADS; + unsigned int warp_tid = ::cuda::ptx::get_sreg_laneid(); + + // Place unpacked digit counters in shared memory + _CCCL_PRAGMA_UNROLL_FULL() + for (int LANE = 0; LANE < LANES_PER_WARP; LANE++) + { + int counter_lane = (LANE * WARPS) + warp_id; + if (counter_lane < COUNTER_LANES) + { + int digit_row = counter_lane << LOG_PACKING_RATIO; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int UNPACKED_COUNTER = 0; UNPACKED_COUNTER < PACKING_RATIO; UNPACKED_COUNTER++) + { + int bin_idx = digit_row + UNPACKED_COUNTER; + + temp_storage.block_counters[warp_tid][bin_idx] = local_counts[LANE][UNPACKED_COUNTER]; + } + } + } + + __syncthreads(); + + // Rake-reduce bin_count reductions + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + bin_count[track] = 0; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < WARP_THREADS; ++i) + { + bin_count[track] += temp_storage.block_counters[i][bin_idx]; + } + } + } + } +}; +} // namespace detail::radix_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_reduce.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_reduce.cuh new file mode 100644 index 00000000..b78a26d7 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_reduce.cuh @@ -0,0 +1,611 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::AgentReduce implements a stateful abstraction of CUDA thread blocks for participating in device-wide reduction. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +namespace detail +{ +// TODO(bgruber): drop in CCCL 4.0 +/** + * Parameterizable tuning policy type for AgentReduce + * @tparam NominalThreadsPerBlock4B Threads per thread block + * @tparam NominalItemsPerThread4B Items per thread (per tile of input) + * @tparam ComputeT Dominant compute type + * @tparam VectorLoadLength Number of items per vectorized load + * @tparam BlockAlgorithm Cooperative block-wide reduction algorithm to use + * @tparam LoadModifier Cache load modifier for reading input elements + */ +template > +struct agent_reduce_policy : ScalingType +{ + /// Number of items per vectorized load + static constexpr int VECTOR_LOAD_LENGTH = VectorLoadLength; + + /// Cooperative block-wide reduction algorithm to use + static constexpr BlockReduceAlgorithm BLOCK_ALGORITHM = BlockAlgorithm; + + /// Cache load modifier for reading input elements + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentReducePolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce") = detail::agent_reduce_policy< + NominalThreadsPerBlock4B, + NominalItemsPerThread4B, + ComputeT, + VectorLoadLength, + BlockAlgorithm, + LoadModifier, + ScalingType>; + +namespace detail +{ +template +struct agent_warp_reduce_policy +{ + /// Number of threads per warp + static constexpr int WARP_THREADS = WarpThreads; + + /// Number of items per vectorized load + static constexpr int VECTOR_LOAD_LENGTH = VectorLoadLength; + + /// Number of threads per block + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + + /// Number of items per thread. When `ComputeT` is `void`, the nominal value is used as-is (no scaling), + /// allowing to pass actual items_per_thread to opt out of the legacy 4B scaling. + static constexpr int ITEMS_PER_THREAD = + ::cuda::std::conditional_t<::cuda::std::is_same_v, + NoScaling<0, NominalItemsPerThread4B>, + MemBoundScaling<0, NominalItemsPerThread4B, ComputeT>>::ITEMS_PER_THREAD; + + /// Cache load modifier for reading input elements + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; + + /// Number of items per tile + constexpr static int ITEMS_PER_TILE = ITEMS_PER_THREAD * WARP_THREADS; + + /// Number of segments per block + constexpr static int SEGMENTS_PER_BLOCK = BLOCK_THREADS / WARP_THREADS; + + static_assert((BLOCK_THREADS % WARP_THREADS) == 0, "Block should be multiple of warp"); +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template +using AgentWarpReducePolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceSegmentedReduce") = detail:: + agent_warp_reduce_policy; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::reduce +{ +/** + * @brief AgentReduceImpl implements a stateful abstraction of CUDA thread blocks + * and warps, for participating in device-wide reduction . + * + * Each thread reduces only the values it loads. If `FIRST_TILE`, this partial + * reduction is stored into `thread_aggregate`. Otherwise it is accumulated + * into `thread_aggregate`. + * + * @tparam AgentReducePolicy + * Parameterized AgentReducePolicy tuning policy type + * + * @tparam InputIteratorT + * Random-access iterator type for input + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam ReductionOp + * Binary reduction operator type having member + * `auto operator()(T &&a, U &&b)` + * + * @tparam AccumT + * The type of intermediate accumulator (according to P2322R6) + * + * @tparam TransformOp + * Unary operator type having member `auto operator()(T &&a)` + * + * @tparam CollectiveReduceT + * Block or Warp reduction type + * + * @tparam NumThreads + * Number of threads participating in the collective reduction + * + * @tparam IsWarpReduction + * Whether or not this is a warp reduction + */ +template +struct AgentReduceImpl +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + /// The input value type + using InputT = it_value_t; + + /// Vector type of InputT for data movement + using VectorT = typename CubVector::Type; + + /// Input iterator wrapper type (for applying cache modifier) + // Wrap the native input pointer with CacheModifiedInputIterator + // or directly use the supplied input iterator type + using WrappedInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + InputIteratorT>; + + /// Constants + static constexpr int ITEMS_PER_THREAD = AgentReducePolicy::ITEMS_PER_THREAD; + static constexpr int TILE_ITEMS = NumThreads * ITEMS_PER_THREAD; + static constexpr int vec_size = ::cuda::std::min(ITEMS_PER_THREAD, AgentReducePolicy::VECTOR_LOAD_LENGTH); + + // Can vectorize according to the policy if the input iterator is a native + // pointer to a primitive type + // TODO(bgruber): we should not check for `is_pointer_v` but `contiguous_iterator` and unwrap it + static constexpr bool ATTEMPT_VECTORIZATION = + (vec_size > 1) && (ITEMS_PER_THREAD % vec_size == 0) + && (::cuda::std::is_pointer_v) + // TODO(bgruber): remove the check for is_primitive in CCCL 4.0 + &&(is_primitive::value || THRUST_NS_QUALIFIER::is_trivially_relocatable_v) + // vectorizing large types leads to regressions again, see https://github.com/NVIDIA/cccl/issues/9761 + // TODO(bgruber): this should be decided by tuning + &&sizeof(InputT) + <= 8; + + static constexpr CacheLoadModifier LOAD_MODIFIER = AgentReducePolicy::LOAD_MODIFIER; + + /// Shared memory type required by this thread block + struct _TempStorage + { + typename CollectiveReduceT::TempStorage reduce; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + _TempStorage& temp_storage; ///< Reference to temp_storage + InputIteratorT d_in; ///< Input data to reduce + WrappedInputIteratorT d_wrapped_in; ///< Wrapped input data to reduce + ReductionOp reduction_op; ///< Binary reduction operator + TransformOp transform_op; ///< Transform operator + unsigned int lane_id; ///< Local thread index inside a Warp or Block + + //--------------------------------------------------------------------- + // Utility + //--------------------------------------------------------------------- + + // Whether the input is aligned with the vector type + template + [[nodiscard]] _CCCL_DEVICE_API static bool IsAligned(Iterator d_in) noexcept + { + if constexpr (AttemptVectorization) + { + return ::cuda::std::is_sufficiently_aligned(d_in); + } + else + { + return false; + } + } + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + + /** + * @brief Constructor + * @param temp_storage Reference to temp_storage + * @param d_in Input data to reduce + * @param reduction_op Binary reduction operator + */ + _CCCL_DEVICE _CCCL_FORCEINLINE AgentReduceImpl( + TempStorage& temp_storage, InputIteratorT d_in, ReductionOp reduction_op, TransformOp transform_op, int lane_id) + : temp_storage(temp_storage.Alias()) + , d_in(d_in) + , d_wrapped_in(d_in) + , reduction_op(reduction_op) + , transform_op(transform_op) + , lane_id(lane_id) + {} + + //--------------------------------------------------------------------- + // Tile consumption + //--------------------------------------------------------------------- + + /** + * @brief Consume a full tile of input + * @tparam IsFirstTile Whether this is a full tile + * @param block_offset The offset the tile to consume + * @param input_is_vector_aligned Whether we can vectorize loads + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeFullTile(AccumT& thread_aggregate, OffsetT block_offset) + { + if constexpr (CanVectorize) + { + // Fabricate a vectorized input iterator + InputT* d_in_unqualified = const_cast(d_in) + block_offset + (lane_id * vec_size); + CacheModifiedInputIterator d_vec_in( + reinterpret_cast(d_in_unqualified)); + + // Load items as vector items + InputT input_items[ITEMS_PER_THREAD]; + VectorT* vec_items = reinterpret_cast(input_items); + + // Alias items as an array of VectorT and load it in striped fashion + static constexpr int words = ITEMS_PER_THREAD / vec_size; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < words; ++i) + { + vec_items[i] = d_vec_in[NumThreads * i]; + } + + // Convert from input type to output type + AccumT items[ITEMS_PER_THREAD]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ITEMS_PER_THREAD; ++i) + { + items[i] = transform_op(input_items[i]); + } + + // Reduce items within each thread stripe + thread_aggregate = + IsFirstTile ? cub::ThreadReduce(items, reduction_op) : cub::ThreadReduce(items, reduction_op, thread_aggregate); + } + else + { + // Scalar path: load items in striped fashion and reduce items within each thread stripe + AccumT items[ITEMS_PER_THREAD]; + load_transform_direct_striped(lane_id, d_wrapped_in + block_offset, items, transform_op); + thread_aggregate = + IsFirstTile ? cub::ThreadReduce(items, reduction_op) : cub::ThreadReduce(items, reduction_op, thread_aggregate); + } + } + + /** + * Consume a partial tile of input + * @tparam IsFirstTile Whether or not this is a full tile + * @param block_offset The offset the tile to consume + * @param valid_items The number of valid items in the tile + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumePartialTile(AccumT& thread_aggregate, OffsetT block_offset, int valid_items) + { + // Partial tile + int thread_offset = lane_id; + + // Read first item + if (IsFirstTile && (thread_offset < valid_items)) + { + thread_aggregate = + transform_op(d_wrapped_in[block_offset + thread_offset]); // NOLINT(bugprone-misplaced-widening-cast) + thread_offset += NumThreads; + } + + // Continue reading items (block-striped) + while (thread_offset < valid_items) + { + InputT item(d_wrapped_in[block_offset + thread_offset]); // NOLINT(bugprone-misplaced-widening-cast) + + thread_aggregate = reduction_op(thread_aggregate, transform_op(item)); + thread_offset += NumThreads; + } + } + + //--------------------------------------------------------------- + // Consume a contiguous segment of tiles + //--------------------------------------------------------------------- + + /** + * @brief Reduce a contiguous segment of input tiles + * @param even_share GridEvenShare descriptor + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE AccumT ConsumeRange(GridEvenShare& even_share) + { + AccumT thread_aggregate{}; + + if (even_share.block_end - even_share.block_offset < TILE_ITEMS) + { + // First tile isn't full (not all threads have valid items) + int valid_items = even_share.block_end - even_share.block_offset; + ConsumePartialTile(thread_aggregate, even_share.block_offset, valid_items); + + // For Warp Reduction, we need to explicitly handle the valid_items, + // whereas for Block Reduction it is implicitly handled + if constexpr (IsWarpReduction) + { + valid_items = (NumThreads <= valid_items) ? NumThreads : valid_items; + } + return CollectiveReduceT(temp_storage.reduce).Reduce(thread_aggregate, reduction_op, valid_items); + } + + // Extracting this into a function saves 8% of generated kernel size by allowing to reuse + // the block reduction below. This also workaround hang in nvcc. + ConsumeFullTileRange(thread_aggregate, even_share); + + // Compute block-wide reduction (all threads have valid items) + return CollectiveReduceT(temp_storage.reduce).Reduce(thread_aggregate, reduction_op); + } + + /** + * @brief Reduce a contiguous segment of input tiles + * @param[in] block_offset Threadblock begin offset (inclusive) + * @param[in] block_end Threadblock end offset (exclusive) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE AccumT ConsumeRange(OffsetT block_offset, OffsetT block_end) + { + GridEvenShare even_share; + even_share.template BlockInit(block_offset, block_end); + + return IsAligned(d_in + block_offset) + ? ConsumeRange(even_share) + : ConsumeRange(even_share); + } + + /** + * Reduce a contiguous segment of input tiles + * @param[in] even_share GridEvenShare descriptor + */ + _CCCL_DEVICE _CCCL_FORCEINLINE AccumT ConsumeTiles(GridEvenShare& even_share) + { + // Initialize GRID_MAPPING_STRIP_MINE even-share descriptor for this thread block + even_share.template BlockInit(); + + return IsAligned(d_in) ? ConsumeRange(even_share) : ConsumeRange(even_share); + } + +private: + /** + * @brief Reduce a contiguous segment of input tiles with more than `TILE_ITEMS` elements + * @param even_share GridEvenShare descriptor + * @param input_is_vector_aligned Whether we can vectorize loads + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeFullTileRange(AccumT& thread_aggregate, GridEvenShare& even_share) + { + // At least one full block + ConsumeFullTile(thread_aggregate, even_share.block_offset); + + if (even_share.block_end - even_share.block_offset < even_share.block_stride) + { + // Exit early to handle offset overflow + return; + } + + even_share.block_offset += even_share.block_stride; + + // Consume subsequent full tiles of input, at least one full tile was processed, so + // `even_share.block_end >= TILE_ITEMS` + while (even_share.block_offset <= even_share.block_end - TILE_ITEMS) + { + ConsumeFullTile(thread_aggregate, even_share.block_offset); + + if (even_share.block_end - even_share.block_offset < even_share.block_stride) + { + // Exit early to handle offset overflow + return; + } + + even_share.block_offset += even_share.block_stride; + } + + // Consume a partially-full tile + if (even_share.block_offset < even_share.block_end) + { + int valid_items = even_share.block_end - even_share.block_offset; + ConsumePartialTile(thread_aggregate, even_share.block_offset, valid_items); + } + } +}; + +/** + * @brief AgentReduce implements a stateful abstraction of CUDA thread blocks + * and warps, for participating in device-wide reduction . + * + * Each thread reduces only the values it loads. If `FIRST_TILE`, this partial + * reduction is stored into `thread_aggregate`. Otherwise it is accumulated + * into `thread_aggregate`. + * + * @tparam AgentReducePolicy + * Parameterized AgentReducePolicy tuning policy type + * + * @tparam InputIteratorT + * Random-access iterator type for input + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam ReductionOp + * Binary reduction operator type having member + * `auto operator()(T &&a, U &&b)` + * + * @tparam AccumT + * The type of intermediate accumulator (according to P2322R6) + * + * @tparam TransformOp + * Unary operator type having member `auto operator()(T &&a)` + */ +template +struct AgentReduce + : AgentReduceImpl, + AgentReducePolicy::BLOCK_THREADS> +{ + using base_t = + AgentReduceImpl, + AgentReducePolicy::BLOCK_THREADS>; + + _CCCL_DEVICE _CCCL_FORCEINLINE AgentReduce( + typename base_t::TempStorage& temp_storage, + InputIteratorT d_in, + ReductionOp reduction_op, + TransformOp transform_op = {}) + : base_t(temp_storage, d_in, reduction_op, transform_op, threadIdx.x) + {} +}; + +/** + * @brief AgentWarpReduce implements a stateful abstraction of CUDA warps, + * for participating in device-wide reduction . + * + * Each thread reduces only the values it loads. If `FIRST_TILE`, this partial + * reduction is stored into `thread_aggregate`. Otherwise it is accumulated + * into `thread_aggregate`. + * + * @tparam AgentReducePolicy + * Parameterized AgentReducePolicy tuning policy type + * + * @tparam InputIteratorT + * Random-access iterator type for input + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam ReductionOp + * Binary reduction operator type having member + * `auto operator()(T &&a, U &&b)` + * + * @tparam AccumT + * The type of intermediate accumulator (according to P2322R6) + * + * @tparam TransformOp + * Unary operator type having member `auto operator()(T &&a)` + */ +template +struct AgentWarpReduce + : AgentReduceImpl, + AgentReducePolicy::WARP_THREADS, + true> +{ + using base_t = + AgentReduceImpl, + AgentReducePolicy::WARP_THREADS, + true>; + + _CCCL_DEVICE _CCCL_FORCEINLINE AgentWarpReduce( + typename base_t::TempStorage& temp_storage, + InputIteratorT d_in, + ReductionOp reduction_op, + TransformOp transform_op = {}) + : base_t(temp_storage, d_in, reduction_op, transform_op, threadIdx.x % AgentReducePolicy::WARP_THREADS) + {} +}; +} // namespace detail::reduce + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_reduce_by_key.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_reduce_by_key.cuh new file mode 100644 index 00000000..e44d6c66 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_reduce_by_key.cuh @@ -0,0 +1,757 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::detail::reduce_by_key::AgentReduceByKey implements a stateful abstraction of CUDA thread blocks for +//! participating in device-wide reduce-value-by-key. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +namespace detail +{ +template > +struct agent_reduce_by_key_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int ITEMS_PER_THREAD = ItemsPerThread; + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + + struct detail + { + using delay_constructor_t = DelayConstructorT; + }; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentReduceByKeyPolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce::ReduceByKey") = detail:: + agent_reduce_by_key_policy; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::reduce_by_key +{ +/** + * @brief AgentReduceByKey implements a stateful abstraction of CUDA thread + * blocks for participating in device-wide reduce-value-by-key + * + * @tparam AgentReduceByKeyPolicyT + * Parameterized AgentReduceByKeyPolicy tuning policy type + * + * @tparam KeysInputIteratorT + * Random-access input iterator type for keys + * + * @tparam UniqueOutputIteratorT + * Random-access output iterator type for keys + * + * @tparam ValuesInputIteratorT + * Random-access input iterator type for values + * + * @tparam AggregatesOutputIteratorT + * Random-access output iterator type for values + * + * @tparam NumRunsOutputIteratorT + * Output iterator type for recording number of items selected + * + * @tparam EqualityOpT + * KeyT equality operator type + * + * @tparam ReductionOpT + * ValueT reduction operator type + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam AccumT + * The type of intermediate accumulator (according to P2322R6) + */ +template +struct AgentReduceByKey +{ + // Whether or not this is a streaming invocation (i.e., multiple kernel invocations over partitions of the input) + static constexpr bool is_streaming_invocation = !::cuda::std::is_same_v; + + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + // The input keys type + using KeyInputT = it_value_t; + + // The output keys type + using KeyOutputT = non_void_value_t; + + // The input values type + using ValueInputT = it_value_t; + + // Tuple type for scanning (pairs accumulated segment-value with + // segment-index) + using OffsetValuePairT = KeyValuePair; + + // Tuple type for pairing keys and values + using KeyValuePairT = KeyValuePair; + + // Tile status descriptor interface type + using ScanTileStateT = ReduceByKeyScanTileState; + + // Guarded inequality functor + template + struct GuardedInequalityWrapper + { + /// Wrapped equality operator + _EqualityOpT op; + + /// Items remaining + int num_remaining; + + /// Constructor + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE GuardedInequalityWrapper(_EqualityOpT op, int num_remaining) + : op(op) + , num_remaining(num_remaining) + {} + + /// Boolean inequality operator, returns (a != b) + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator()(const T& a, const T& b, int idx) const + { + if (idx < num_remaining) + { + return !op(a, b); // In bounds + } + + // Return true if first out-of-bounds item, false otherwise + return (idx == num_remaining); + } + }; + + // Constants + static constexpr int BLOCK_THREADS = AgentReduceByKeyPolicyT::BLOCK_THREADS; + static constexpr int ITEMS_PER_THREAD = AgentReduceByKeyPolicyT::ITEMS_PER_THREAD; + static constexpr int TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD; + static constexpr int TWO_PHASE_SCATTER = (ITEMS_PER_THREAD > 1); + + // Cache-modified Input iterator wrapper type (for applying cache modifier) + // for keys Wrap the native input pointer with + // CacheModifiedValuesInputIterator or directly use the supplied input + // iterator type + using WrappedKeysInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + KeysInputIteratorT>; + + // Cache-modified Input iterator wrapper type (for applying cache modifier) + // for values Wrap the native input pointer with + // CacheModifiedValuesInputIterator or directly use the supplied input + // iterator type + using WrappedValuesInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + ValuesInputIteratorT>; + + // Cache-modified Input iterator wrapper type (for applying cache modifier) + // for fixup values Wrap the native input pointer with + // CacheModifiedValuesInputIterator or directly use the supplied input + // iterator type + using WrappedFixupInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + AggregatesOutputIteratorT>; + + // Reduce-value-by-segment scan operator + using ReduceBySegmentOpT = ReduceBySegmentOp; + + // Parameterized BlockLoad type for keys + using BlockLoadKeysT = + BlockLoad; + + // Parameterized BlockLoad type for values + using BlockLoadValuesT = BlockLoad; + + // Parameterized BlockDiscontinuity type for keys + using BlockDiscontinuityKeys = BlockDiscontinuity; + + // Parameterized BlockScan type + using BlockScanT = BlockScan; + + // Callback type for obtaining tile prefix during block scan + using DelayConstructorT = typename AgentReduceByKeyPolicyT::detail::delay_constructor_t; + using TilePrefixCallbackOpT = + TilePrefixCallbackOp; + + // Key and value exchange types + using KeyExchangeT = KeyOutputT[TILE_ITEMS + 1]; + using ValueExchangeT = AccumT[TILE_ITEMS + 1]; + + // Shared memory type for this thread block + union _TempStorage + { + struct ScanStorage + { + // Smem needed for tile scanning + typename BlockScanT::TempStorage scan; + + // Smem needed for cooperative prefix callback + typename TilePrefixCallbackOpT::TempStorage prefix; + + // Smem needed for discontinuity detection + typename BlockDiscontinuityKeys::TempStorage discontinuity; + } scan_storage; + + // Smem needed for loading keys + typename BlockLoadKeysT::TempStorage load_keys; + + // Smem needed for loading values + typename BlockLoadValuesT::TempStorage load_values; + + // Smem needed for compacting key value pairs(allows non POD items in this + // union) + Uninitialized raw_exchange; + }; + + // Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + /// Reference to temp_storage + _TempStorage& temp_storage; + + /// Input keys + WrappedKeysInputIteratorT d_keys_in; + + /// Unique output keys + UniqueOutputIteratorT d_unique_out; + + /// Input values + WrappedValuesInputIteratorT d_values_in; + + /// Output value aggregates + AggregatesOutputIteratorT d_aggregates_out; + + /// Output pointer for total number of segments identified + NumRunsOutputIteratorT d_num_runs_out; + + /// KeyT equality operator + EqualityOpT equality_op; + + /// Reduction operator + ReductionOpT reduction_op; + + /// Reduce-by-segment scan operator + ReduceBySegmentOpT scan_op; + + /// Streaming context providing context about this partition for streaming invocations + StreamingContextT streaming_context; + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + + /** + * @param temp_storage + * Reference to temp_storage + * + * @param d_keys_in + * Input keys + * + * @param d_unique_out + * Unique output keys + * + * @param d_values_in + * Input values + * + * @param d_aggregates_out + * Output value aggregates + * + * @param d_num_runs_out + * Output pointer for total number of segments identified + * + * @param equality_op + * KeyT equality operator + * + * @param reduction_op + * ValueT reduction operator + * + * @param streaming_context + * Streaming context providing context about this partition for streaming invocations + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE AgentReduceByKey( + TempStorage& temp_storage, + KeysInputIteratorT d_keys_in, + UniqueOutputIteratorT d_unique_out, + ValuesInputIteratorT d_values_in, + AggregatesOutputIteratorT d_aggregates_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + ReductionOpT reduction_op, + StreamingContext streaming_context) + : temp_storage(temp_storage.Alias()) + , d_keys_in(d_keys_in) + , d_unique_out(d_unique_out + streaming_context.num_uniques()) + , d_values_in(d_values_in) + , d_aggregates_out(d_aggregates_out + streaming_context.num_uniques()) + , d_num_runs_out(d_num_runs_out) + , equality_op(equality_op) + , reduction_op(reduction_op) + , scan_op(reduction_op) + , streaming_context(streaming_context) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE AgentReduceByKey( + TempStorage& temp_storage, + KeysInputIteratorT d_keys_in, + UniqueOutputIteratorT d_unique_out, + ValuesInputIteratorT d_values_in, + AggregatesOutputIteratorT d_aggregates_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + ReductionOpT reduction_op, + NullType streaming_context) + : temp_storage(temp_storage.Alias()) + , d_keys_in(d_keys_in) + , d_unique_out(d_unique_out) + , d_values_in(d_values_in) + , d_aggregates_out(d_aggregates_out) + , d_num_runs_out(d_num_runs_out) + , equality_op(equality_op) + , reduction_op(reduction_op) + , scan_op(reduction_op) + , streaming_context(streaming_context) + {} + + //--------------------------------------------------------------------- + // Scatter utility methods + //--------------------------------------------------------------------- + + /** + * Directly scatter flagged items to output offsets + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterDirect( + KeyValuePairT (&scatter_items)[ITEMS_PER_THREAD], + OffsetT (&segment_flags)[ITEMS_PER_THREAD], + OffsetT (&segment_indices)[ITEMS_PER_THREAD]) + { + // Scatter flagged keys and values + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + if (segment_flags[ITEM]) + { + d_unique_out[segment_indices[ITEM]] = scatter_items[ITEM].key; + d_aggregates_out[segment_indices[ITEM]] = scatter_items[ITEM].value; + } + } + } + + /** + * 2-phase scatter flagged items to output offsets + * + * The exclusive scan causes each head flag to be paired with the previous + * value aggregate: the scatter offsets must be decremented for value + * aggregates + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterTwoPhase( + KeyValuePairT (&scatter_items)[ITEMS_PER_THREAD], + OffsetT (&segment_flags)[ITEMS_PER_THREAD], + OffsetT (&segment_indices)[ITEMS_PER_THREAD], + OffsetT num_tile_segments, + OffsetT num_tile_segments_prefix) + { + __syncthreads(); + + // Compact and scatter pairs + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + if (segment_flags[ITEM]) + { + temp_storage.raw_exchange.Alias()[segment_indices[ITEM] - num_tile_segments_prefix] = scatter_items[ITEM]; + } + } + + __syncthreads(); + + for (int item = static_cast(threadIdx.x); item < num_tile_segments; item += BLOCK_THREADS) + { + KeyValuePairT pair = temp_storage.raw_exchange.Alias()[item]; + d_unique_out[num_tile_segments_prefix + item] = pair.key; // NOLINT(bugprone-misplaced-widening-cast) + d_aggregates_out[num_tile_segments_prefix + item] = pair.value; // NOLINT(bugprone-misplaced-widening-cast) + } + } + + /** + * Scatter flagged items + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void Scatter( + KeyValuePairT (&scatter_items)[ITEMS_PER_THREAD], + OffsetT (&segment_flags)[ITEMS_PER_THREAD], + OffsetT (&segment_indices)[ITEMS_PER_THREAD], + OffsetT num_tile_segments, + OffsetT num_tile_segments_prefix) + { + // Do a one-phase scatter if (a) two-phase is disabled or (b) the average + // number of selected items per thread is less than one + if (TWO_PHASE_SCATTER && (num_tile_segments > BLOCK_THREADS)) + { + ScatterTwoPhase(scatter_items, segment_flags, segment_indices, num_tile_segments, num_tile_segments_prefix); + } + else + { + ScatterDirect(scatter_items, segment_flags, segment_indices); + } + } + + //--------------------------------------------------------------------- + // Cooperatively scan a device-wide sequence of tiles with other CTAs + //--------------------------------------------------------------------- + + /** + * @brief Process a tile of input (dynamic chained scan) + * + * @tparam IS_LAST_TILE + * Whether the current tile is the last tile + * + * @param num_remaining + * Number of global input items remaining (including this tile) + * + * @param tile_idx + * Tile index + * + * @param tile_offset + * Tile offset + * + * @param tile_state + * Global tile state descriptor + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeTile(OffsetT num_remaining, int tile_idx, OffsetT tile_offset, ScanTileStateT& tile_state) + { + // Tile keys + KeyOutputT keys[ITEMS_PER_THREAD]; + + // Tile keys shuffled up + KeyOutputT prev_keys[ITEMS_PER_THREAD]; + + // Tile values + AccumT values[ITEMS_PER_THREAD]; + + // Segment head flags + OffsetT head_flags[ITEMS_PER_THREAD]; + + // Segment indices + OffsetT segment_indices[ITEMS_PER_THREAD]; + + // Zipped values and segment flags|indices + OffsetValuePairT scan_items[ITEMS_PER_THREAD]; + + // Zipped key value pairs for scattering + KeyValuePairT scatter_items[ITEMS_PER_THREAD]; + + // Load keys + if (IS_LAST_TILE) + { + BlockLoadKeysT(temp_storage.load_keys).Load(d_keys_in + tile_offset, keys, num_remaining); + } + else + { + BlockLoadKeysT(temp_storage.load_keys).Load(d_keys_in + tile_offset, keys); + } + + // Load tile predecessor key in first thread + KeyOutputT tile_predecessor; + if (threadIdx.x == 0) + { + // if (tile_idx == 0) + // first tile gets repeat of first item (thus first item will not + // be flagged as a head) + // else + // Subsequent tiles get last key from previous tile + if constexpr (is_streaming_invocation) + { + tile_predecessor = (tile_idx == 0) ? streaming_context.predecessor_key() : d_keys_in[tile_offset - 1]; + } + else + { + tile_predecessor = (tile_idx == 0) ? keys[0] : d_keys_in[tile_offset - 1]; + } + } + + __syncthreads(); + + // Load values + if (IS_LAST_TILE) + { + BlockLoadValuesT(temp_storage.load_values).Load(d_values_in + tile_offset, values, num_remaining); + } + else + { + BlockLoadValuesT(temp_storage.load_values).Load(d_values_in + tile_offset, values); + } + + __syncthreads(); + + // Initialize head-flags and shuffle up the previous keys + if (IS_LAST_TILE) + { + // Use custom flag operator to additionally flag the first out-of-bounds item + GuardedInequalityWrapper flag_op(equality_op, num_remaining); + BlockDiscontinuityKeys(temp_storage.scan_storage.discontinuity) + .FlagHeads(head_flags, keys, prev_keys, flag_op, tile_predecessor); + } + else + { + InequalityWrapper flag_op(equality_op); + BlockDiscontinuityKeys(temp_storage.scan_storage.discontinuity) + .FlagHeads(head_flags, keys, prev_keys, flag_op, tile_predecessor); + } + + // Reset head-flag on the very first item to make sure we don't start a new run for data where + // (key[0] == key[0]) is false (e.g., when key[0] is NaN) + if constexpr (is_streaming_invocation) + { + if (streaming_context.is_first_partition() && threadIdx.x == 0 && tile_idx == 0) + { + head_flags[0] = 0; + } + } + else + { + if (threadIdx.x == 0 && tile_idx == 0) + { + head_flags[0] = 0; + } + } + + // Zip values and head flags + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + scan_items[ITEM].value = values[ITEM]; + scan_items[ITEM].key = head_flags[ITEM]; + } + + // Perform exclusive tile scan + // Inclusive block-wide scan aggregate + OffsetValuePairT block_aggregate; + + // Number of segments prior to this tile + OffsetT num_segments_prefix; + + // The tile prefix folded with block_aggregate + OffsetValuePairT total_aggregate; + + if (tile_idx == 0) + { + // Scan first tile + // First partition does not need to account for preceding partitions + if constexpr (is_streaming_invocation) + { + if (streaming_context.is_first_partition()) + { + BlockScanT(temp_storage.scan_storage.scan).ExclusiveScan(scan_items, scan_items, scan_op, block_aggregate); + num_segments_prefix = 0; + total_aggregate = block_aggregate; + } + // Subsequent partitions need to account for preceding partitions + else + { + auto init_value = OffsetValuePairT{0, streaming_context.prefix()}; + BlockScanT(temp_storage.scan_storage.scan) + .ExclusiveScan(scan_items, scan_items, init_value, scan_op, block_aggregate); + num_segments_prefix = 0; + // note, block_aggregate does not include the prefix + block_aggregate = scan_op(init_value, block_aggregate); + total_aggregate = block_aggregate; + } + } + else + { + BlockScanT(temp_storage.scan_storage.scan).ExclusiveScan(scan_items, scan_items, scan_op, block_aggregate); + num_segments_prefix = 0; + total_aggregate = block_aggregate; + } + + // Update tile status if there are successor tiles + if ((!IS_LAST_TILE) && (threadIdx.x == 0)) + { + tile_state.SetInclusive(0, block_aggregate); + } + } + else + { + // Scan non-first tile + TilePrefixCallbackOpT prefix_op(tile_state, temp_storage.scan_storage.prefix, scan_op, tile_idx); + BlockScanT(temp_storage.scan_storage.scan).ExclusiveScan(scan_items, scan_items, scan_op, prefix_op); + + block_aggregate = prefix_op.GetBlockAggregate(); + num_segments_prefix = prefix_op.GetExclusivePrefix().key; + total_aggregate = prefix_op.GetInclusivePrefix(); + } + + // Rezip scatter items and segment indices + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + scatter_items[ITEM].key = prev_keys[ITEM]; + scatter_items[ITEM].value = scan_items[ITEM].value; + segment_indices[ITEM] = scan_items[ITEM].key; + } + + // At this point, each flagged segment head has: + // - The key for the previous segment + // - The reduced value from the previous segment + // - The segment index for the reduced value + + // Scatter flagged keys and values + OffsetT num_tile_segments = block_aggregate.key; + Scatter(scatter_items, head_flags, segment_indices, num_tile_segments, num_segments_prefix); + + // Last thread in last tile will output final count (and last pair, if necessary) + if ((IS_LAST_TILE) && (threadIdx.x == BLOCK_THREADS - 1)) + { + OffsetT num_segments = num_segments_prefix + num_tile_segments; + + // If the last tile is a full tile, we need to write out the run ending with the last item + // If this was not a full tile, we already have flagged the head of one-past-the-last-item + if (num_remaining == TILE_ITEMS) + { + if constexpr (is_streaming_invocation) + { + if (streaming_context.is_last_partition()) + { + d_unique_out[num_segments] = keys[ITEMS_PER_THREAD - 1]; + d_aggregates_out[num_segments] = total_aggregate.value; + num_segments++; + } + else + { + // Write the prefix aggregate of this partition as context for the subsequent partition + streaming_context.write_prefix(total_aggregate.value); + } + } + else + { + d_unique_out[num_segments] = keys[ITEMS_PER_THREAD - 1]; + d_aggregates_out[num_segments] = total_aggregate.value; + num_segments++; + } + } + + if constexpr (is_streaming_invocation) + { + // Add the number of unique items in this partition to the global aggregate + auto total_uniques = streaming_context.add_num_uniques(num_segments); + + // If this is the last partition, write out the number of unique items + if (streaming_context.is_last_partition()) + { + *d_num_runs_out = total_uniques; + } + } + else + { + *d_num_runs_out = num_segments; + } + } + } + + /** + * @brief Scan tiles of items as part of a dynamic chained scan + * + * @param num_items + * Total number of input items + * + * @param tile_state + * Global tile state descriptor + * + * @param start_tile + * The starting tile for the current grid + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeRange(OffsetT num_items, ScanTileStateT& tile_state, int start_tile) + { + // Blocks are launched in increasing order, so just assign one tile per + // block + + // Current tile index + int tile_idx = static_cast(start_tile + blockIdx.x); + + // Global offset for the current tile + OffsetT tile_offset = OffsetT(TILE_ITEMS) * tile_idx; + + // Remaining items (including this tile) + OffsetT num_remaining = num_items - tile_offset; + + if (num_remaining > TILE_ITEMS) + { + // Not last tile + ConsumeTile(num_remaining, tile_idx, tile_offset, tile_state); + } + else if (num_remaining > 0) + { + // Last tile + ConsumeTile(num_remaining, tile_idx, tile_offset, tile_state); + } + } +}; +} // namespace detail::reduce_by_key + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_rle.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_rle.cuh new file mode 100644 index 00000000..22c48e7c --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_rle.cuh @@ -0,0 +1,1072 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * cub::AgentRle implements a stateful abstraction of CUDA thread blocks for participating in device-wide + * run-length-encode. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +namespace detail +{ +template > +struct agent_rle_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int ITEMS_PER_THREAD = ItemsPerThread; + static constexpr bool STORE_WARP_TIME_SLICING = StoreWarpTimeSlicing; + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + + struct detail + { + using delay_constructor_t = DelayConstructorT; + }; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentRlePolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceRunLengthEncode") = + detail::agent_rle_policy; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::rle +{ +/** + * @brief AgentRle implements a stateful abstraction of CUDA thread blocks for participating in device-wide + * run-length-encode + * + * @tparam AgentRlePolicyT + * Parameterized AgentRlePolicyT tuning policy type + * + * @tparam InputIteratorT + * Random-access input iterator type for data + * + * @tparam OffsetsOutputIteratorT + * Random-access output iterator type for offset values + * + * @tparam LengthsOutputIteratorT + * Random-access output iterator type for length values + * + * @tparam EqualityOpT + * T equality operator type + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam StreamingContextT + * Type providing information about the partition for streaming invocations. NullType if not a streaming invocation. + */ +template +struct AgentRle +{ + // Whether or not this is a streaming invocation (i.e., multiple kernel invocations over partitions of the input) + static constexpr bool is_streaming_invocation = !::cuda::std::is_same_v; + + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + /// The input value type + using T = cub::detail::it_value_t; + + /// The lengths output value type + using LengthT = cub::detail::non_void_value_t; + + /// Tuple type for scanning (pairs run-length and run-index) + using LengthOffsetPair = KeyValuePair; + + /// Tile status descriptor interface type + using ScanTileStateT = ReduceByKeyScanTileState; + + // Constants + static constexpr int WARP_THREADS = warp_threads; + static constexpr int BLOCK_THREADS = AgentRlePolicyT::BLOCK_THREADS; + static constexpr int ITEMS_PER_THREAD = AgentRlePolicyT::ITEMS_PER_THREAD; + static constexpr int WARP_ITEMS = WARP_THREADS * ITEMS_PER_THREAD; + static constexpr int TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD; + static constexpr int WARPS = (BLOCK_THREADS + WARP_THREADS - 1) / WARP_THREADS; + + /// Whether or not to sync after loading data + static constexpr bool SYNC_AFTER_LOAD = (AgentRlePolicyT::LOAD_ALGORITHM != BLOCK_LOAD_DIRECT); + + /// Whether or not only one warp's worth of shared memory should be allocated and time-sliced + /// among block-warps during any store-related data transpositions (versus each warp having + /// its own storage) + static constexpr bool STORE_WARP_TIME_SLICING = AgentRlePolicyT::STORE_WARP_TIME_SLICING; + static constexpr int ACTIVE_EXCHANGE_WARPS = (STORE_WARP_TIME_SLICING) ? 1 : WARPS; + + /** + * Special operator that signals all out-of-bounds items are not equal to everything else, + * forcing both (1) the last item to be tail-flagged and (2) all oob items to be marked + * trivial. + */ + template + struct OobInequalityOp + { + OffsetT num_remaining; + EqualityOpT equality_op; + + _CCCL_DEVICE _CCCL_FORCEINLINE OobInequalityOp(OffsetT num_remaining, EqualityOpT equality_op) + : num_remaining(num_remaining) + , equality_op(equality_op) + {} + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator()(T first, T second, Index idx) + { + if (!LAST_TILE || (idx < num_remaining)) + { + return !equality_op(first, second); + } + else + { + return true; + } + } + }; + + // Cache-modified Input iterator wrapper type (for applying cache modifier) for data + // Wrap the native input pointer with CacheModifiedVLengthnputIterator + // Directly use the supplied input iterator type + using WrappedInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + InputIteratorT>; + + // Parameterized BlockLoad type for data + using BlockLoadT = + BlockLoad; + + // Parameterized BlockDiscontinuity type for data + using BlockDiscontinuityT = BlockDiscontinuity; + + // Parameterized WarpScan type + using WarpScanPairs = WarpScan; + + // Reduce-length-by-run scan operator + using ReduceBySegmentOpT = ReduceBySegmentOp<::cuda::std::plus<>>; + + // Callback type for obtaining tile prefix during block scan + using DelayConstructorT = typename AgentRlePolicyT::detail::delay_constructor_t; + using TilePrefixCallbackOpT = + TilePrefixCallbackOp; + + // Warp exchange types + using WarpExchangePairs = WarpExchange; + + using WarpExchangePairsStorage = + ::cuda::std::_If; + + using WarpExchangeOffsets = WarpExchange; + using WarpExchangeLengths = WarpExchange; + + using WarpAggregates = LengthOffsetPair[WARPS]; + + // Shared memory type for this thread block + struct _TempStorage + { + // Aliasable storage layout + union Aliasable + { + struct ScanStorage + { + // Smem needed for discontinuity detection + typename BlockDiscontinuityT::TempStorage discontinuity; + + // Smem needed for warp-synchronous scans + typename WarpScanPairs::TempStorage warp_scan[WARPS]; + + // Smem needed for sharing warp-wide aggregates + Uninitialized warp_aggregates; + + // Smem needed for cooperative prefix callback + typename TilePrefixCallbackOpT::TempStorage prefix; + } scan_storage; + + // Smem needed for input loading + typename BlockLoadT::TempStorage load; + + // Aliasable layout needed for two-phase scatter + union ScatterAliasable + { + unsigned long long align; + WarpExchangePairsStorage exchange_pairs[ACTIVE_EXCHANGE_WARPS]; + typename WarpExchangeOffsets::TempStorage exchange_offsets[ACTIVE_EXCHANGE_WARPS]; + typename WarpExchangeLengths::TempStorage exchange_lengths[ACTIVE_EXCHANGE_WARPS]; + } scatter_aliasable; + + } aliasable; + + OffsetT tile_idx; // Shared tile index + LengthOffsetPair tile_inclusive; // Inclusive tile prefix + LengthOffsetPair tile_exclusive; // Exclusive tile prefix + }; + + // Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + _TempStorage& temp_storage; ///< Reference to temp_storage + + WrappedInputIteratorT d_in; ///< Pointer to input sequence of data items + OffsetsOutputIteratorT d_offsets_out; ///< Input run offsets + LengthsOutputIteratorT d_lengths_out; ///< Output run lengths + + EqualityOpT equality_op; ///< T equality operator + ReduceBySegmentOpT scan_op; ///< Reduce-length-by-flag scan operator + OffsetT num_items; ///< Total number of input items + StreamingContextT streaming_context; ///< Context providing information about this partition for streaming invocations + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + + /** + * @param[in] temp_storage + * Reference to temp_storage + * + * @param[in] d_in + * Pointer to input sequence of data items + * + * @param[out] d_offsets_out + * Pointer to output sequence of run offsets + * + * @param[out] d_lengths_out + * Pointer to output sequence of run lengths + * + * @param[in] equality_op + * Equality operator + * + * @param[in] num_items + * Total number of input items + * + * @param streaming_context + * Streaming context providing context about this partition for streaming invocations + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE AgentRle( + TempStorage& temp_storage, + InputIteratorT d_in, + OffsetsOutputIteratorT d_offsets_out, + LengthsOutputIteratorT d_lengths_out, + EqualityOpT equality_op, + OffsetT num_items, + StreamingContext streaming_context) + : temp_storage(temp_storage.Alias()) + , d_in(d_in) + , d_offsets_out(d_offsets_out) + , d_lengths_out(d_lengths_out) + , equality_op(equality_op) + , scan_op(::cuda::std::plus<>{}) + , num_items(num_items) + , streaming_context(streaming_context) + {} + + //--------------------------------------------------------------------- + // Utility methods for initializing the selections + //--------------------------------------------------------------------- + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InitializeSelections( + OffsetT tile_offset, + OffsetT num_remaining, + T (&items)[ITEMS_PER_THREAD], + LengthOffsetPair (&lengths_and_num_runs)[ITEMS_PER_THREAD]) + { + bool head_flags[ITEMS_PER_THREAD]; + bool tail_flags[ITEMS_PER_THREAD]; + + OobInequalityOp inequality_op(num_remaining, equality_op); + + if (FIRST_TILE && LAST_TILE) + { + // First-and-last-tile always head-flags the first item and tail-flags the last item + + BlockDiscontinuityT(temp_storage.aliasable.scan_storage.discontinuity) + .FlagHeadsAndTails(head_flags, tail_flags, items, inequality_op); + } + else if (FIRST_TILE) + { + // First-tile always head-flags the first item + + // Get the first item from the next tile + T tile_successor_item; + if (threadIdx.x == BLOCK_THREADS - 1) + { + tile_successor_item = d_in[tile_offset + TILE_ITEMS]; // NOLINT(bugprone-misplaced-widening-cast) + } + + BlockDiscontinuityT(temp_storage.aliasable.scan_storage.discontinuity) + .FlagHeadsAndTails(head_flags, tail_flags, tile_successor_item, items, inequality_op); + } + else if (LAST_TILE) + { + // Last-tile always flags the last item + + // Get the last item from the previous tile + T tile_predecessor_item; + if (threadIdx.x == 0) + { + tile_predecessor_item = d_in[tile_offset - 1]; + } + + BlockDiscontinuityT(temp_storage.aliasable.scan_storage.discontinuity) + .FlagHeadsAndTails(head_flags, tile_predecessor_item, tail_flags, items, inequality_op); + } + else + { + // Get the first item from the next tile + T tile_successor_item; + if (threadIdx.x == BLOCK_THREADS - 1) + { + tile_successor_item = d_in[tile_offset + TILE_ITEMS]; // NOLINT(bugprone-misplaced-widening-cast) + } + + // Get the last item from the previous tile + T tile_predecessor_item; + if (threadIdx.x == 0) + { + tile_predecessor_item = d_in[tile_offset - 1]; + } + + BlockDiscontinuityT(temp_storage.aliasable.scan_storage.discontinuity) + .FlagHeadsAndTails(head_flags, tile_predecessor_item, tail_flags, tile_successor_item, items, inequality_op); + } + + // Zip counts and runs + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + // input output + // items [ 0 0 0 1 2 3 3 ] + // heads [ 1 0 0 1 1 1 0 ] + // tails [ 0 0 1 1 1 0 1 ] + // key [ 1 0 0 0 0 1 0 ] head && !tail - heads of non-trivial (length > 1) runs + // value [ 1 1 1 0 0 1 1 ] !head || !tail - elements of non-trivial runs + lengths_and_num_runs[ITEM].key = head_flags[ITEM] && (!tail_flags[ITEM]); + lengths_and_num_runs[ITEM].value = ((!head_flags[ITEM]) || (!tail_flags[ITEM])); + } + } + + //--------------------------------------------------------------------- + // Scan utility methods + //--------------------------------------------------------------------- + + /** + * Scan of allocations + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void WarpScanAllocations( + LengthOffsetPair& tile_aggregate, + LengthOffsetPair& warp_aggregate, + LengthOffsetPair& warp_exclusive_in_tile, + LengthOffsetPair& thread_exclusive_in_warp, + LengthOffsetPair (&lengths_and_num_runs)[ITEMS_PER_THREAD]) + { + // Perform warpscans + unsigned int warp_id = ((WARPS == 1) ? 0 : threadIdx.x / WARP_THREADS); + int lane_id = static_cast(::cuda::ptx::get_sreg_laneid()); + + LengthOffsetPair identity; + identity.key = 0; + identity.value = 0; + + LengthOffsetPair thread_inclusive; + + // `thread_exclusive_in_warp.key`: + // number of non-trivial runs starts in previous threads + // `thread_exclusive_in_warp.val`: + // number of items in the last non-trivial run in previous threads + + // `thread_aggregate.key`: + // number of non-trivial runs starts in this thread + // `thread_aggregate.val`: + // number of items in the last non-trivial run in this thread + LengthOffsetPair thread_aggregate = cub::ThreadReduce(lengths_and_num_runs, scan_op); + WarpScanPairs(temp_storage.aliasable.scan_storage.warp_scan[warp_id]) + .Scan(thread_aggregate, thread_inclusive, thread_exclusive_in_warp, identity, scan_op); + + // `thread_inclusive.key`: + // number of non-trivial runs starts in this and previous warp threads + // `thread_inclusive.val`: + // number of items in the last non-trivial run in this or previous warp threads + + // Last lane in each warp shares its warp-aggregate + if (lane_id == WARP_THREADS - 1) + { + // `temp_storage.aliasable.scan_storage.warp_aggregates[warp_id].key`: + // number of non-trivial runs starts in this warp + // `temp_storage.aliasable.scan_storage.warp_aggregates[warp_id].val`: + // number of items in the last non-trivial run in this warp + temp_storage.aliasable.scan_storage.warp_aggregates.Alias()[warp_id] = thread_inclusive; + } + + __syncthreads(); + + // Accumulate total selected and the warp-wide prefix + + // `warp_exclusive_in_tile.key`: + // number of non-trivial runs starts in previous warps + // `warp_exclusive_in_tile.val`: + // number of items in the last non-trivial run in previous warps + warp_exclusive_in_tile = identity; + warp_aggregate = temp_storage.aliasable.scan_storage.warp_aggregates.Alias()[warp_id]; + + // `tile_aggregate.key`: + // number of non-trivial runs starts in this CTA + // `tile_aggregate.val`: + // number of items in the last non-trivial run in this CTA + tile_aggregate = temp_storage.aliasable.scan_storage.warp_aggregates.Alias()[0]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int WARP = 1; WARP < WARPS; ++WARP) + { + if (warp_id == WARP) + { + warp_exclusive_in_tile = tile_aggregate; + } + + tile_aggregate = scan_op(tile_aggregate, temp_storage.aliasable.scan_storage.warp_aggregates.Alias()[WARP]); + } + + // Ensure all threads have read warp aggregates before temp_storage is repurposed in the + // subsequent scatter stage + __syncthreads(); + } + + //--------------------------------------------------------------------- + // Utility methods for scattering selections + //--------------------------------------------------------------------- + + /** + * Two-phase scatter, specialized for warp time-slicing + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterTwoPhase( + OffsetT tile_num_runs_exclusive_in_global, + OffsetT warp_num_runs_aggregate, + OffsetT warp_num_runs_exclusive_in_tile, + OffsetT (&thread_num_runs_exclusive_in_warp)[ITEMS_PER_THREAD], + LengthOffsetPair (&lengths_and_offsets)[ITEMS_PER_THREAD], + ::cuda::std::true_type is_warp_time_slice) + { + unsigned int warp_id = ((WARPS == 1) ? 0 : threadIdx.x / WARP_THREADS); + int lane_id = static_cast(::cuda::ptx::get_sreg_laneid()); + + // Locally compact items within the warp (first warp) + if (warp_id == 0) + { + WarpExchangePairs(temp_storage.aliasable.scatter_aliasable.exchange_pairs[0]) + .ScatterToStriped(lengths_and_offsets, thread_num_runs_exclusive_in_warp); + } + + // Locally compact items within the warp (remaining warps) + _CCCL_PRAGMA_UNROLL_FULL() + for (int SLICE = 1; SLICE < WARPS; ++SLICE) + { + __syncthreads(); + + if (warp_id == SLICE) + { + WarpExchangePairs(temp_storage.aliasable.scatter_aliasable.exchange_pairs[0]) + .ScatterToStriped(lengths_and_offsets, thread_num_runs_exclusive_in_warp); + } + } + + // Global scatter + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ITEM++) + { + // warp_num_runs_aggregate - number of non-trivial runs starts in current warp + if ((ITEM * WARP_THREADS) < warp_num_runs_aggregate - lane_id) + { + OffsetT item_offset = + tile_num_runs_exclusive_in_global + warp_num_runs_exclusive_in_tile + (ITEM * WARP_THREADS) + lane_id; + + // Scatter offset + if constexpr (is_streaming_invocation) + { + d_offsets_out[streaming_context.num_uniques() + item_offset] = + (streaming_context.base_offset() + lengths_and_offsets[ITEM].key); + + // Scatter length if not the first (global) length + if (streaming_context.num_uniques() + item_offset > 0) + { + d_lengths_out[streaming_context.num_uniques() + item_offset - 1] = lengths_and_offsets[ITEM].value; + } + } + else + { + d_offsets_out[item_offset] = lengths_and_offsets[ITEM].key; + + // Scatter length if not the first (global) length + if ((ITEM != 0) || (item_offset > 0)) + { + d_lengths_out[item_offset - 1] = lengths_and_offsets[ITEM].value; + } + } + } + } + } + + /** + * Two-phase scatter + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterTwoPhase( + OffsetT tile_num_runs_exclusive_in_global, + OffsetT warp_num_runs_aggregate, + OffsetT warp_num_runs_exclusive_in_tile, + OffsetT (&thread_num_runs_exclusive_in_warp)[ITEMS_PER_THREAD], + LengthOffsetPair (&lengths_and_offsets)[ITEMS_PER_THREAD], + ::cuda::std::false_type is_warp_time_slice) + { + unsigned int warp_id = ((WARPS == 1) ? 0 : threadIdx.x / WARP_THREADS); + int lane_id = static_cast(::cuda::ptx::get_sreg_laneid()); + + // Unzip + OffsetT run_offsets[ITEMS_PER_THREAD]; + LengthT run_lengths[ITEMS_PER_THREAD]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ITEM++) + { + run_offsets[ITEM] = lengths_and_offsets[ITEM].key; + run_lengths[ITEM] = lengths_and_offsets[ITEM].value; + } + + WarpExchangeOffsets(temp_storage.aliasable.scatter_aliasable.exchange_offsets[warp_id]) + .ScatterToStriped(run_offsets, thread_num_runs_exclusive_in_warp); + + __syncwarp(0xffffffff); + + WarpExchangeLengths(temp_storage.aliasable.scatter_aliasable.exchange_lengths[warp_id]) + .ScatterToStriped(run_lengths, thread_num_runs_exclusive_in_warp); + + // Global scatter + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ITEM++) + { + if ((ITEM * WARP_THREADS) + lane_id < warp_num_runs_aggregate) + { + OffsetT item_offset = + tile_num_runs_exclusive_in_global + warp_num_runs_exclusive_in_tile + (ITEM * WARP_THREADS) + lane_id; + + // Scatter offset + if constexpr (is_streaming_invocation) + { + d_offsets_out[streaming_context.num_uniques() + item_offset] = + (streaming_context.base_offset() + run_offsets[ITEM]); + // Scatter length if not the first (global) length + if ((ITEM != 0) || (streaming_context.num_uniques() + item_offset > 0)) + { + d_lengths_out[streaming_context.num_uniques() + item_offset - 1] = run_lengths[ITEM]; + } + } + else + { + d_offsets_out[item_offset] = run_offsets[ITEM]; + // Scatter length if not the first (global) length + if ((ITEM != 0) || (item_offset > 0)) + { + d_lengths_out[item_offset - 1] = run_lengths[ITEM]; + } + } + } + } + } + + /** + * Direct scatter + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterDirect( + OffsetT tile_num_runs_exclusive_in_global, + OffsetT warp_num_runs_aggregate, + OffsetT warp_num_runs_exclusive_in_tile, + OffsetT (&thread_num_runs_exclusive_in_warp)[ITEMS_PER_THREAD], + LengthOffsetPair (&lengths_and_offsets)[ITEMS_PER_THREAD]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + if (thread_num_runs_exclusive_in_warp[ITEM] < warp_num_runs_aggregate) + { + OffsetT item_offset = + tile_num_runs_exclusive_in_global + warp_num_runs_exclusive_in_tile + thread_num_runs_exclusive_in_warp[ITEM]; + + // Scatter offset + if constexpr (is_streaming_invocation) + { + // For streaming invocations, we need to add the base offset of the partition + d_offsets_out[streaming_context.num_uniques() + item_offset] = + (streaming_context.base_offset() + lengths_and_offsets[ITEM].key); + + // Scatter length if not the first (global) length + if (streaming_context.num_uniques() + item_offset > 0) + { + d_lengths_out[streaming_context.num_uniques() + item_offset - 1] = lengths_and_offsets[ITEM].value; + } + } + else + { + d_offsets_out[item_offset] = lengths_and_offsets[ITEM].key; + + // Scatter length if not the first (global) length + if (item_offset > 0) + { + d_lengths_out[item_offset - 1] = lengths_and_offsets[ITEM].value; + } + } + } + } + } + + /** + * Scatter + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void Scatter( + OffsetT tile_num_runs_aggregate, + OffsetT tile_num_runs_exclusive_in_global, + OffsetT warp_num_runs_aggregate, + OffsetT warp_num_runs_exclusive_in_tile, + OffsetT (&thread_num_runs_exclusive_in_warp)[ITEMS_PER_THREAD], + LengthOffsetPair (&lengths_and_offsets)[ITEMS_PER_THREAD]) + { + if ((ITEMS_PER_THREAD == 1) || (tile_num_runs_aggregate < BLOCK_THREADS)) + { + // Direct scatter if the warp has any items + if (warp_num_runs_aggregate) + { + ScatterDirect(tile_num_runs_exclusive_in_global, + warp_num_runs_aggregate, + warp_num_runs_exclusive_in_tile, + thread_num_runs_exclusive_in_warp, + lengths_and_offsets); + } + } + else + { + // Scatter two phase + ScatterTwoPhase( + tile_num_runs_exclusive_in_global, + warp_num_runs_aggregate, + warp_num_runs_exclusive_in_tile, + thread_num_runs_exclusive_in_warp, + lengths_and_offsets, + bool_constant_v); + } + } + + //--------------------------------------------------------------------- + // Cooperatively scan a device-wide sequence of tiles with other CTAs + //--------------------------------------------------------------------- + + /** + * @brief Process a tile of input (dynamic chained scan) + * + * @param num_items + * Total number of global input items + * + * @param num_remaining + * Number of global input items remaining (including this tile) + * + * @param tile_idx + * Tile index + * + * @param tile_offset + * Tile offset + * + * @param &tile_status + * Global list of tile status + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE LengthOffsetPair + ConsumeTile(OffsetT num_items, OffsetT num_remaining, int tile_idx, OffsetT tile_offset, ScanTileStateT& tile_status) + { + if (tile_idx == 0) + { + // First tile + + // Load items + T items[ITEMS_PER_THREAD]; + if (LAST_TILE) + { + BlockLoadT(temp_storage.aliasable.load).Load(d_in + tile_offset, items, num_remaining, T()); + } + else + { + BlockLoadT(temp_storage.aliasable.load).Load(d_in + tile_offset, items); + } + + if (SYNC_AFTER_LOAD) + { + __syncthreads(); + } + + // Set flags + LengthOffsetPair lengths_and_num_runs[ITEMS_PER_THREAD]; + + if constexpr (is_streaming_invocation) + { + if (streaming_context.first_partition) + { + if (streaming_context.last_partition) + { + InitializeSelections(tile_offset, num_remaining, items, lengths_and_num_runs); + } + else + { + InitializeSelections(tile_offset, num_remaining, items, lengths_and_num_runs); + } + } + else + { + if (streaming_context.last_partition) + { + InitializeSelections(tile_offset, num_remaining, items, lengths_and_num_runs); + } + else + { + InitializeSelections(tile_offset, num_remaining, items, lengths_and_num_runs); + } + } + } + else + { + InitializeSelections(tile_offset, num_remaining, items, lengths_and_num_runs); + } + + // Exclusive scan of lengths and runs + LengthOffsetPair tile_aggregate; + LengthOffsetPair warp_aggregate; + LengthOffsetPair warp_exclusive_in_tile; + LengthOffsetPair thread_exclusive_in_warp; + + if constexpr (is_streaming_invocation) + { + // If this is a streaming invocation, we need to incorporate the run-length of the previous partition's last run + if (!streaming_context.first_partition && threadIdx.x == 0) + { + lengths_and_num_runs[0].value += streaming_context.prefix(); + } + } + + WarpScanAllocations( + tile_aggregate, warp_aggregate, warp_exclusive_in_tile, thread_exclusive_in_warp, lengths_and_num_runs); + + // Update tile status if this is not the last tile + if (!LAST_TILE && (threadIdx.x == 0)) + { + tile_status.SetInclusive(0, tile_aggregate); + } + + // Update thread_exclusive_in_warp to fold in warp run-length + if (thread_exclusive_in_warp.key == 0) + { + // If there are no non-trivial runs starts in the previous warp threads, then + // `thread_exclusive_in_warp.val` denotes the number of items in the last + // non-trivial run of the previous CTA threads, so the better name for it is + // `thread_exclusive_in_tile`. + thread_exclusive_in_warp.value += warp_exclusive_in_tile.value; + } + + LengthOffsetPair lengths_and_offsets[ITEMS_PER_THREAD]; + OffsetT thread_num_runs_exclusive_in_warp[ITEMS_PER_THREAD]; + LengthOffsetPair lengths_and_num_runs2[ITEMS_PER_THREAD]; + + // Downsweep scan through lengths_and_num_runs + detail::ThreadScanExclusive(lengths_and_num_runs, lengths_and_num_runs2, scan_op, thread_exclusive_in_warp); + + // Zip + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ITEM++) + { + lengths_and_offsets[ITEM].value = lengths_and_num_runs2[ITEM].value; + lengths_and_offsets[ITEM].key = tile_offset + (threadIdx.x * ITEMS_PER_THREAD) + ITEM; + thread_num_runs_exclusive_in_warp[ITEM] = + (lengths_and_num_runs[ITEM].key) ? lengths_and_num_runs2[ITEM].key : // keep + WARP_THREADS * ITEMS_PER_THREAD; // discard + } + + OffsetT tile_num_runs_aggregate = tile_aggregate.key; + OffsetT tile_num_runs_exclusive_in_global = 0; + OffsetT warp_num_runs_aggregate = warp_aggregate.key; + OffsetT warp_num_runs_exclusive_in_tile = warp_exclusive_in_tile.key; + + // Scatter + Scatter(tile_num_runs_aggregate, + tile_num_runs_exclusive_in_global, + warp_num_runs_aggregate, + warp_num_runs_exclusive_in_tile, + thread_num_runs_exclusive_in_warp, + lengths_and_offsets); + + // Return running total (inclusive of this tile) + return tile_aggregate; + } + else + { + // Not first tile + + // Load items + T items[ITEMS_PER_THREAD]; + if (LAST_TILE) + { + BlockLoadT(temp_storage.aliasable.load).Load(d_in + tile_offset, items, num_remaining, T()); + } + else + { + BlockLoadT(temp_storage.aliasable.load).Load(d_in + tile_offset, items); + } + + if (SYNC_AFTER_LOAD) + { + __syncthreads(); + } + + // Set flags + LengthOffsetPair lengths_and_num_runs[ITEMS_PER_THREAD]; + + if constexpr (is_streaming_invocation) + { + if (streaming_context.last_partition) + { + InitializeSelections(tile_offset, num_remaining, items, lengths_and_num_runs); + } + else + { + InitializeSelections(tile_offset, num_remaining, items, lengths_and_num_runs); + } + } + else + { + InitializeSelections(tile_offset, num_remaining, items, lengths_and_num_runs); + } + + // Exclusive scan of lengths and runs + LengthOffsetPair tile_aggregate; + LengthOffsetPair warp_aggregate; + LengthOffsetPair warp_exclusive_in_tile; + LengthOffsetPair thread_exclusive_in_warp; + + WarpScanAllocations( + tile_aggregate, warp_aggregate, warp_exclusive_in_tile, thread_exclusive_in_warp, lengths_and_num_runs); + + // First warp computes tile prefix in lane 0 + TilePrefixCallbackOpT prefix_op( + tile_status, temp_storage.aliasable.scan_storage.prefix, ::cuda::std::plus<>{}, tile_idx); + unsigned int warp_id = ((WARPS == 1) ? 0 : threadIdx.x / WARP_THREADS); + if (warp_id == 0) + { + prefix_op(tile_aggregate); + if (threadIdx.x == 0) + { + temp_storage.tile_exclusive = prefix_op.exclusive_prefix; + } + } + + __syncthreads(); + + LengthOffsetPair tile_exclusive_in_global = temp_storage.tile_exclusive; + + // Update thread_exclusive_in_warp to fold in warp and tile run-lengths + LengthOffsetPair thread_exclusive = scan_op(tile_exclusive_in_global, warp_exclusive_in_tile); + if (thread_exclusive_in_warp.key == 0) + { + // If there are no non-trivial runs starts in the previous warp threads, then + // `thread_exclusive_in_warp.val` denotes the number of items in the last + // non-trivial run of the previous grid threads, so the better name for it is + // `thread_exclusive_in_grid`. + thread_exclusive_in_warp.value += thread_exclusive.value; + } + + // Downsweep scan through lengths_and_num_runs + + // `lengths_and_num_runs2.key`: + // number of non-trivial runs starts in previous grid threads + // `lengths_and_num_runs2.val`: + // number of items in the last non-trivial run in previous grid threads + LengthOffsetPair lengths_and_num_runs2[ITEMS_PER_THREAD]; + + // `lengths_and_offsets.key`: + // offset to the item in the input sequence + // `lengths_and_offsets.val`: + // number of items in the last non-trivial run in previous grid threads + LengthOffsetPair lengths_and_offsets[ITEMS_PER_THREAD]; + OffsetT thread_num_runs_exclusive_in_warp[ITEMS_PER_THREAD]; + + detail::ThreadScanExclusive(lengths_and_num_runs, lengths_and_num_runs2, scan_op, thread_exclusive_in_warp); + + // Zip + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ITEM++) + { + lengths_and_offsets[ITEM].value = lengths_and_num_runs2[ITEM].value; + lengths_and_offsets[ITEM].key = tile_offset + (threadIdx.x * ITEMS_PER_THREAD) + ITEM; + thread_num_runs_exclusive_in_warp[ITEM] = + (lengths_and_num_runs[ITEM].key) ? lengths_and_num_runs2[ITEM].key : // keep + WARP_THREADS * ITEMS_PER_THREAD; // discard + } + + OffsetT tile_num_runs_aggregate = tile_aggregate.key; + OffsetT tile_num_runs_exclusive_in_global = tile_exclusive_in_global.key; + OffsetT warp_num_runs_aggregate = warp_aggregate.key; + OffsetT warp_num_runs_exclusive_in_tile = warp_exclusive_in_tile.key; + + // Scatter + Scatter(tile_num_runs_aggregate, + tile_num_runs_exclusive_in_global, + warp_num_runs_aggregate, + warp_num_runs_exclusive_in_tile, + thread_num_runs_exclusive_in_warp, + lengths_and_offsets); + + // Return running total (inclusive of this tile) + return prefix_op.inclusive_prefix; + } + } + + /** + * @brief Scan tiles of items as part of a dynamic chained scan + * + * @param num_tiles + * Total number of input tiles + * + * @param tile_status + * Global list of tile status + * + * @param d_num_runs_out + * Output pointer for total number of runs identified + * + * @tparam NumRunsIteratorT + * Output iterator type for recording number of items selected + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeRange(int num_tiles, ScanTileStateT& tile_status, NumRunsIteratorT d_num_runs_out) + { + // Blocks are launched in increasing order, so just assign one tile per block + int tile_idx = static_cast((blockIdx.x * gridDim.y) + blockIdx.y); // Current tile index + OffsetT tile_offset = static_cast(tile_idx) * static_cast(TILE_ITEMS); + OffsetT num_remaining = num_items - tile_offset; // Remaining items (including this tile) + + if (tile_idx < num_tiles - 1) + { + // Not the last tile (full) + ConsumeTile(num_items, num_remaining, tile_idx, tile_offset, tile_status); + } + else if (num_remaining > 0) + { + // The last tile (possibly partially-full) + LengthOffsetPair running_total = ConsumeTile(num_items, num_remaining, tile_idx, tile_offset, tile_status); + + if (threadIdx.x == 0) + { + if constexpr (is_streaming_invocation) + { + // Add the number of unique items in this partition to the global aggregate + auto total_uniques = streaming_context.add_num_uniques(running_total.key); + + // If this is the last partition, write out the number of unique items + if (streaming_context.last_partition) + { + // Output the total number of items selected + *d_num_runs_out = total_uniques; + + // The inclusive prefix contains accumulated length reduction for the last run + if (running_total.key + streaming_context.num_uniques() > 0) + { + d_lengths_out[streaming_context.num_uniques() + running_total.key - 1] = running_total.value; + } + } + + if (!streaming_context.last_partition) + { + // Write the run-length of this partition as context for the subsequent partition + streaming_context.write_prefix(running_total.value); + } + } + else + { + // Output the total number of items selected + *d_num_runs_out = running_total.key; + + // The inclusive prefix contains accumulated length reduction for the last run + if (running_total.key > 0) + { + d_lengths_out[running_total.key - 1] = running_total.value; + } + } + } + } + } +}; +} // namespace detail::rle + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_scan.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_scan.cuh new file mode 100644 index 00000000..1bda6f3d --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_scan.cuh @@ -0,0 +1,567 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * @brief cub::AgentScan implements a stateful abstraction of CUDA thread blocks + * for participating in device-wide prefix scan. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +// TODO(bgruber): remove when C++20 is the minimum, since then we can pass policy values as NTTPs +template , + typename DelayConstructorT = detail::default_delay_constructor_t> +struct agent_scan_policy : ScalingType +{ + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; + static constexpr BlockStoreAlgorithm STORE_ALGORITHM = StoreAlgorithm; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + + struct detail + { + using delay_constructor_t = DelayConstructorT; + }; +}; +} // namespace detail + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +/** + * @brief Parameterizable tuning policy type for AgentScan + * + * @tparam NominalThreadsPerBlock4B + * Threads per thread block + * + * @tparam NominalItemsPerThread4B + * Items per thread (per tile of input) + * + * @tparam ComputeT + * Dominant compute type + * + * @tparam LoadAlgorithm + * The BlockLoad algorithm to use + * + * @tparam LoadModifier + * Cache load modifier for reading input elements + * + * @tparam StoreAlgorithm + * The BlockStore algorithm to use + * + * @tparam ScanAlgorithm + * The BlockScan algorithm to use + * + * @tparam DelayConstructorT + * Implementation detail, do not specify directly, requirements on the + * content of this type are subject to breaking change. + */ +//! Deprecated [Since 3.5] +template , + typename DelayConstructorT = detail::default_delay_constructor_t> +using AgentScanPolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") = detail::agent_scan_policy< + NominalThreadsPerBlock4B, + NominalItemsPerThread4B, + ComputeT, + LoadAlgorithm, + LoadModifier, + StoreAlgorithm, + ScanAlgorithm, + ScalingType, + DelayConstructorT>; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::scan +{ +/** + * @brief AgentScan implements a stateful abstraction of CUDA thread blocks for + * participating in device-wide prefix scan. + * @tparam AgentScanPolicyT + * Parameterized AgentScanPolicyT tuning policy type + * + * @tparam InputIteratorT + * Random-access input iterator type + * + * @tparam OutputIteratorT + * Random-access output iterator type + * + * @tparam ScanOpT + * Scan functor type + * + * @tparam InitValueT + * The init_value element for ScanOpT type (cub::NullType for inclusive scan) + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam AccumT + * The type of intermediate accumulator (according to P2322R6) + */ +template +struct AgentScan +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + // The input value type + using InputT = cub::detail::it_value_t; + + // Tile status descriptor interface type + using ScanTileStateT = ScanTileState; + + // Input iterator wrapper type (for applying cache modifier) + // Wrap the native input pointer with CacheModifiedInputIterator + // or directly use the supplied input iterator type + using WrappedInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + InputIteratorT>; + + // Inclusive scan if no init_value type is provided + static constexpr bool HAS_INIT = !::cuda::std::is_same_v; + static constexpr bool IS_INCLUSIVE = ForceInclusive || !HAS_INIT; // We are relying on either initial value not being + // `NullType` or the ForceInclusive tag to be true + // for inclusive scan to get picked up. + static constexpr int BLOCK_THREADS = AgentScanPolicyT::BLOCK_THREADS; + static constexpr int ITEMS_PER_THREAD = AgentScanPolicyT::ITEMS_PER_THREAD; + static constexpr int TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD; + + // Parameterized BlockLoad type + using BlockLoadT = + BlockLoad; + + // Parameterized BlockStore type + using BlockStoreT = + BlockStore; + + // Parameterized BlockScan type + using BlockScanT = BlockScan; + + // Callback type for obtaining tile prefix during block scan + using DelayConstructorT = typename AgentScanPolicyT::detail::delay_constructor_t; + using TilePrefixCallbackOpT = + TilePrefixCallbackOp; + + // Stateful BlockScan prefix callback type for managing a running total while + // scanning consecutive tiles + using RunningPrefixCallbackOp = BlockScanRunningPrefixOp; + + // Shared memory type for this thread block + union _TempStorage + { + // Smem needed for tile loading + typename BlockLoadT::TempStorage load; + + // Smem needed for tile storing + typename BlockStoreT::TempStorage store; + + struct ScanStorage + { + // Smem needed for cooperative prefix callback + typename TilePrefixCallbackOpT::TempStorage prefix; + + // Smem needed for tile scanning + typename BlockScanT::TempStorage scan; + } scan_storage; + }; + + // Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + _TempStorage& temp_storage; ///< Reference to temp_storage + WrappedInputIteratorT d_in; ///< Input data + OutputIteratorT d_out; ///< Output data + ScanOpT scan_op; ///< Binary scan operator + InitValueT init_value; ///< The init_value element for ScanOpT + + //--------------------------------------------------------------------- + // Block scan utility methods + //--------------------------------------------------------------------- + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ScanFirstTile(AccumT (&items)[ITEMS_PER_THREAD], InitValueT init_value, ScanOpT scan_op, AccumT& block_aggregate) + { + BlockScanT blockScan(temp_storage.scan_storage.scan); + if constexpr (Inclusive) + { + if constexpr (HAS_INIT) + { + blockScan.InclusiveScan(items, items, init_value, scan_op, block_aggregate); + block_aggregate = scan_op(init_value, block_aggregate); + } + else + { + blockScan.InclusiveScan(items, items, scan_op, block_aggregate); + } + } + else + { + blockScan.ExclusiveScan(items, items, init_value, scan_op, block_aggregate); + block_aggregate = scan_op(init_value, block_aggregate); + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ScanSubsequentTile(AccumT (&items)[ITEMS_PER_THREAD], ScanOpT scan_op, PrefixCallback& prefix_op) + { + BlockScanT blockScan(temp_storage.scan_storage.scan); + if constexpr (Inclusive) + { + blockScan.InclusiveScan(items, items, scan_op, prefix_op); + } + else + { + blockScan.ExclusiveScan(items, items, scan_op, prefix_op); + } + } + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + + /** + * @param temp_storage + * Reference to temp_storage + * + * @param d_in + * Input data + * + * @param d_out + * Output data + * + * @param scan_op + * Binary scan operator + * + * @param init_value + * Initial value to seed the exclusive scan + */ + _CCCL_DEVICE _CCCL_FORCEINLINE AgentScan( + TempStorage& temp_storage, InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op, InitValueT init_value) + : temp_storage(temp_storage.Alias()) + , d_in(d_in) + , d_out(d_out) + , scan_op(scan_op) + , init_value(init_value) + {} + + //--------------------------------------------------------------------- + // Cooperatively scan a device-wide sequence of tiles with other CTAs + //--------------------------------------------------------------------- + + /** + * Process a tile of input (dynamic chained scan) + * @tparam IS_LAST_TILE + * Whether the current tile is the last tile + * + * @param num_remaining + * Number of global input items remaining (including this tile) + * + * @param tile_idx + * Tile index + * + * @param tile_offset + * Tile offset + * + * @param tile_state + * Global tile state descriptor + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeTile(OffsetT num_remaining, int tile_idx, OffsetT tile_offset, ScanTileStateT& tile_state) + { + // Load items + AccumT items[ITEMS_PER_THREAD]; + + if constexpr (IS_LAST_TILE) + { + // Fill last element with the first element because collectives are + // not suffix guarded. + BlockLoadT(temp_storage.load).Load(d_in + tile_offset, items, num_remaining, *(d_in + tile_offset)); + } + else + { + BlockLoadT(temp_storage.load).Load(d_in + tile_offset, items); + } + + __syncthreads(); + + // Perform tile scan + if (tile_idx == 0) + { + // Scan first tile + AccumT block_aggregate; + ScanFirstTile(items, init_value, scan_op, block_aggregate); + + if ((!IS_LAST_TILE) && (threadIdx.x == 0)) + { + tile_state.SetInclusive(0, block_aggregate); + } + } + else + { + // Scan non-first tile + TilePrefixCallbackOpT prefix_op(tile_state, temp_storage.scan_storage.prefix, scan_op, tile_idx); + ScanSubsequentTile(items, scan_op, prefix_op); + } + + __syncthreads(); + + if constexpr (UsePDL) + { + _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); // omitting makes almost no difference in cub.bench.scan.exclusive.sum.base + } + + // Store items + if constexpr (IS_LAST_TILE) + { + BlockStoreT(temp_storage.store).Store(d_out + tile_offset, items, num_remaining); + } + else + { + BlockStoreT(temp_storage.store).Store(d_out + tile_offset, items); + } + } + + /** + * @brief Scan tiles of items as part of a dynamic chained scan + * + * @param num_items + * Total number of input items + * + * @param tile_state + * Global tile state descriptor + * + * @param start_tile + * The starting tile for the current grid + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeRange(OffsetT num_items, ScanTileStateT& tile_state, int start_tile) + { + // Blocks are launched in increasing order, so just assign one tile per + // block + + // Current tile index + int tile_idx = static_cast(start_tile + blockIdx.x); + + // Global offset for the current tile + OffsetT tile_offset = OffsetT(TILE_ITEMS) * tile_idx; + + // Remaining items (including this tile) + OffsetT num_remaining = num_items - tile_offset; + + if (num_remaining > TILE_ITEMS) + { + // Not last tile + ConsumeTile(num_remaining, tile_idx, tile_offset, tile_state); + } + else if (num_remaining > 0) + { + // Last tile + ConsumeTile(num_remaining, tile_idx, tile_offset, tile_state); + } + } + + //--------------------------------------------------------------------------- + // Scan an sequence of consecutive tiles (independent of other thread blocks) + //--------------------------------------------------------------------------- + + /** + * @brief Process a tile of input + * + * @param tile_offset + * Tile offset + * + * @param prefix_op + * Running prefix operator + * + * @param valid_items + * Number of valid items in the tile + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeTile(OffsetT tile_offset, RunningPrefixCallbackOp& prefix_op, int valid_items = TILE_ITEMS) + { + // Load items + AccumT items[ITEMS_PER_THREAD]; + + if constexpr (IS_LAST_TILE) + { + // Fill last element with the first element because collectives are + // not suffix guarded. + BlockLoadT(temp_storage.load).Load(d_in + tile_offset, items, valid_items, *(d_in + tile_offset)); + } + else + { + BlockLoadT(temp_storage.load).Load(d_in + tile_offset, items); + } + + __syncthreads(); + + // Block scan + if constexpr (IS_FIRST_TILE) + { + AccumT block_aggregate; + ScanFirstTile(items, init_value, scan_op, block_aggregate); + prefix_op.running_total = block_aggregate; + } + else + { + ScanSubsequentTile(items, scan_op, prefix_op); + } + + __syncthreads(); + + // Store items + if constexpr (IS_LAST_TILE) + { + BlockStoreT(temp_storage.store).Store(d_out + tile_offset, items, valid_items); + } + else + { + BlockStoreT(temp_storage.store).Store(d_out + tile_offset, items); + } + } + + /** + * @brief Scan a consecutive share of input tiles + * + * @param[in] range_offset + * Threadblock begin offset (inclusive) + * + * @param[in] range_end + * Threadblock end offset (exclusive) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeRange(OffsetT range_offset, OffsetT range_end) + { + BlockScanRunningPrefixOp prefix_op(scan_op); + + if (range_offset + TILE_ITEMS <= range_end) + { + // Consume first tile of input (full) + ConsumeTile(range_offset, prefix_op); + range_offset += TILE_ITEMS; + + // Consume subsequent full tiles of input + while (range_offset + TILE_ITEMS <= range_end) + { + ConsumeTile(range_offset, prefix_op); + range_offset += TILE_ITEMS; + } + + // Consume a partially-full tile + if (range_offset < range_end) + { + int valid_items = range_end - range_offset; + ConsumeTile(range_offset, prefix_op, valid_items); + } + } + else + { + // Consume the first tile of input (partially-full) + int valid_items = range_end - range_offset; + ConsumeTile(range_offset, prefix_op, valid_items); + } + } + + /** + * @brief Scan a consecutive share of input tiles, seeded with the + * specified prefix value + * @param[in] range_offset + * Threadblock begin offset (inclusive) + * + * @param[in] range_end + * Threadblock end offset (exclusive) + * + * @param[in] prefix + * The prefix to apply to the scan segment + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeRange(OffsetT range_offset, OffsetT range_end, AccumT prefix) + { + BlockScanRunningPrefixOp prefix_op(prefix, scan_op); + + // Consume full tiles of input + while (range_offset + TILE_ITEMS <= range_end) + { + ConsumeTile(range_offset, prefix_op); + range_offset += TILE_ITEMS; + } + + // Consume a partially-full tile + if (range_offset < range_end) + { + int valid_items = range_end - range_offset; + ConsumeTile(range_offset, prefix_op, valid_items); + } + } +}; +} // namespace detail::scan + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_scan_by_key.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_scan_by_key.cuh new file mode 100644 index 00000000..d77a4aa2 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_scan_by_key.cuh @@ -0,0 +1,464 @@ +// SPDX-FileCopyrightText: Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * @brief AgentScanByKey implements a stateful abstraction of CUDA thread blocks + * for participating in device-wide prefix scan by key. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +namespace detail +{ +// TODO(bgruber): remove this when C++20 is the minimum, since then we can pass policy values as NTTP +template > +struct agent_scan_by_key_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int ITEMS_PER_THREAD = ItemsPerThread; + + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + static constexpr BlockStoreAlgorithm STORE_ALGORITHM = StoreAlgorithm; + + struct detail + { + using delay_constructor_t = DelayConstructorT; + }; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentScanByKeyPolicy + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScanByKey") = detail::agent_scan_by_key_policy< + ThreadsPerBlock, + ItemsPerThread, + LoadAlgorithm, + LoadModifier, + ScanAlgorithm, + StoreAlgorithm, + DelayConstructorT>; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::scan_by_key +{ +/** + * @brief AgentScanByKey implements a stateful abstraction of CUDA thread + * blocks for participating in device-wide prefix scan by key. + * + * @tparam AgentScanByKeyPolicyT + * Parameterized AgentScanPolicyT tuning policy type + * + * @tparam KeysInputIteratorT + * Random-access input iterator type + * + * @tparam ValuesInputIteratorT + * Random-access input iterator type + * + * @tparam ValuesOutputIteratorT + * Random-access output iterator type + * + * @tparam EqualityOp + * Equality functor type + * + * @tparam ScanOpT + * Scan functor type + * + * @tparam InitValueT + * The init_value element for ScanOpT type (cub::NullType for inclusive scan) + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam AccumT + * The type of intermediate accumulator (according to P2322R6) + */ +template +struct AgentScanByKey +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + using KeyT = it_value_t; + using InputT = it_value_t; + using FlagValuePairT = KeyValuePair; + using ReduceBySegmentOpT = ScanBySegmentOp; + + using ScanTileStateT = ReduceByKeyScanTileState; + + // Constants + // Inclusive scan if no init_value type is provided + static constexpr int IS_INCLUSIVE = ::cuda::std::is_same_v; + static constexpr int BLOCK_THREADS = AgentScanByKeyPolicyT::BLOCK_THREADS; + static constexpr int ITEMS_PER_THREAD = AgentScanByKeyPolicyT::ITEMS_PER_THREAD; + static constexpr int ITEMS_PER_TILE = BLOCK_THREADS * ITEMS_PER_THREAD; + + using WrappedKeysInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + KeysInputIteratorT>; + + using WrappedValuesInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + ValuesInputIteratorT>; + + using BlockLoadKeysT = BlockLoad; + + using BlockLoadValuesT = BlockLoad; + + using BlockStoreValuesT = BlockStore; + + using BlockDiscontinuityKeysT = BlockDiscontinuity; + + using DelayConstructorT = typename AgentScanByKeyPolicyT::detail::delay_constructor_t; + using TilePrefixCallbackT = + TilePrefixCallbackOp; + + using BlockScanT = BlockScan; + + union TempStorage_ + { + struct ScanStorage + { + typename BlockScanT::TempStorage scan; + typename TilePrefixCallbackT::TempStorage prefix; + typename BlockDiscontinuityKeysT::TempStorage discontinuity; + } scan_storage; + + typename BlockLoadKeysT::TempStorage load_keys; + typename BlockLoadValuesT::TempStorage load_values; + typename BlockStoreValuesT::TempStorage store_values; + }; + + struct TempStorage : cub::Uninitialized + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + TempStorage_& storage; + WrappedKeysInputIteratorT d_keys_in; + KeyT* d_keys_prev_in; + WrappedValuesInputIteratorT d_values_in; + ValuesOutputIteratorT d_values_out; + InequalityWrapper inequality_op; + ScanOpT scan_op; + ReduceBySegmentOpT pair_scan_op; + InitValueT init_value; + + //--------------------------------------------------------------------- + // Block scan utility methods (first tile) + //--------------------------------------------------------------------- + + // Exclusive scan specialization + _CCCL_DEVICE _CCCL_FORCEINLINE void + ScanTile(FlagValuePairT (&scan_items)[ITEMS_PER_THREAD], + FlagValuePairT& tile_aggregate, + ::cuda::std::false_type /* is_inclusive */) + { + BlockScanT(storage.scan_storage.scan).ExclusiveScan(scan_items, scan_items, pair_scan_op, tile_aggregate); + } + + // Inclusive scan specialization + _CCCL_DEVICE _CCCL_FORCEINLINE void + ScanTile(FlagValuePairT (&scan_items)[ITEMS_PER_THREAD], + FlagValuePairT& tile_aggregate, + ::cuda::std::true_type /* is_inclusive */) + { + BlockScanT(storage.scan_storage.scan).InclusiveScan(scan_items, scan_items, pair_scan_op, tile_aggregate); + } + + //--------------------------------------------------------------------- + // Block scan utility methods (subsequent tiles) + //--------------------------------------------------------------------- + + // Exclusive scan specialization (with prefix from predecessors) + _CCCL_DEVICE _CCCL_FORCEINLINE void ScanTile( + FlagValuePairT (&scan_items)[ITEMS_PER_THREAD], + FlagValuePairT& tile_aggregate, + TilePrefixCallbackT& prefix_op, + ::cuda::std::false_type /* is_inclusive */) + { + BlockScanT(storage.scan_storage.scan).ExclusiveScan(scan_items, scan_items, pair_scan_op, prefix_op); + tile_aggregate = prefix_op.GetBlockAggregate(); + } + + // Inclusive scan specialization (with prefix from predecessors) + _CCCL_DEVICE _CCCL_FORCEINLINE void ScanTile( + FlagValuePairT (&scan_items)[ITEMS_PER_THREAD], + FlagValuePairT& tile_aggregate, + TilePrefixCallbackT& prefix_op, + ::cuda::std::true_type /* is_inclusive */) + { + BlockScanT(storage.scan_storage.scan).InclusiveScan(scan_items, scan_items, pair_scan_op, prefix_op); + tile_aggregate = prefix_op.GetBlockAggregate(); + } + + //--------------------------------------------------------------------- + // Zip utility methods + //--------------------------------------------------------------------- + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ZipValuesAndFlags( + OffsetT num_remaining, + AccumT (&values)[ITEMS_PER_THREAD], + OffsetT (&segment_flags)[ITEMS_PER_THREAD], + FlagValuePairT (&scan_items)[ITEMS_PER_THREAD]) + { + // Zip values and segment_flags + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + // Set segment_flags for first out-of-bounds item, zero for others + if (IS_LAST_TILE && OffsetT(threadIdx.x * ITEMS_PER_THREAD) + ITEM == num_remaining) + { + segment_flags[ITEM] = 1; + } + + scan_items[ITEM].value = values[ITEM]; + scan_items[ITEM].key = segment_flags[ITEM]; + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + UnzipValues(AccumT (&values)[ITEMS_PER_THREAD], FlagValuePairT (&scan_items)[ITEMS_PER_THREAD]) + { + // Unzip values and segment_flags + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + values[ITEM] = scan_items[ITEM].value; + } + } + + template , ::cuda::std::enable_if_t = 0> + _CCCL_DEVICE _CCCL_FORCEINLINE void + AddInitToScan(AccumT (&items)[ITEMS_PER_THREAD], OffsetT (&flags)[ITEMS_PER_THREAD]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + items[ITEM] = flags[ITEM] ? init_value : scan_op(init_value, items[ITEM]); + } + } + + template , ::cuda::std::enable_if_t = 0> + _CCCL_DEVICE _CCCL_FORCEINLINE void + AddInitToScan(AccumT (& /*items*/)[ITEMS_PER_THREAD], OffsetT (& /*flags*/)[ITEMS_PER_THREAD]) + {} + + //--------------------------------------------------------------------- + // Cooperatively scan a device-wide sequence of tiles with other CTAs + //--------------------------------------------------------------------- + + // Process a tile of input (dynamic chained scan) + // + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeTile(OffsetT /*num_items*/, OffsetT num_remaining, int tile_idx, OffsetT tile_base, ScanTileStateT& tile_state) + { + // Load items + KeyT keys[ITEMS_PER_THREAD]; + AccumT values[ITEMS_PER_THREAD]; + OffsetT segment_flags[ITEMS_PER_THREAD]; + FlagValuePairT scan_items[ITEMS_PER_THREAD]; + + if (IS_LAST_TILE) + { + // Fill last element with the first element + // because collectives are not suffix guarded + BlockLoadKeysT(storage.load_keys).Load(d_keys_in + tile_base, keys, num_remaining, *(d_keys_in + tile_base)); + } + else + { + BlockLoadKeysT(storage.load_keys).Load(d_keys_in + tile_base, keys); + } + + __syncthreads(); + + if (IS_LAST_TILE) + { + // Fill last element with the first element + // because collectives are not suffix guarded + BlockLoadValuesT(storage.load_values) + .Load(d_values_in + tile_base, values, num_remaining, *(d_values_in + tile_base)); + } + else + { + BlockLoadValuesT(storage.load_values).Load(d_values_in + tile_base, values); + } + + __syncthreads(); + + // first tile + if (tile_idx == 0) + { + BlockDiscontinuityKeysT(storage.scan_storage.discontinuity).FlagHeads(segment_flags, keys, inequality_op); + + // Zip values and segment_flags + ZipValuesAndFlags(num_remaining, values, segment_flags, scan_items); + + // Exclusive scan of values and segment_flags + FlagValuePairT tile_aggregate; + ScanTile(scan_items, tile_aggregate, bool_constant_v); + + if (threadIdx.x == 0) + { + if (!IS_LAST_TILE) + { + tile_state.SetInclusive(0, tile_aggregate); + } + + scan_items[0].key = 0; + } + } + else + { + KeyT tile_pred_key = (threadIdx.x == 0) ? d_keys_prev_in[tile_idx] : KeyT(); + + BlockDiscontinuityKeysT(storage.scan_storage.discontinuity) + .FlagHeads(segment_flags, keys, inequality_op, tile_pred_key); + + // Zip values and segment_flags + ZipValuesAndFlags(num_remaining, values, segment_flags, scan_items); + + FlagValuePairT tile_aggregate; + TilePrefixCallbackT prefix_op(tile_state, storage.scan_storage.prefix, pair_scan_op, tile_idx); + ScanTile(scan_items, tile_aggregate, prefix_op, bool_constant_v); + } + + __syncthreads(); + + UnzipValues(values, scan_items); + + AddInitToScan(values, segment_flags); + + // Store items + if (IS_LAST_TILE) + { + BlockStoreValuesT(storage.store_values).Store(d_values_out + tile_base, values, num_remaining); + } + else + { + BlockStoreValuesT(storage.store_values).Store(d_values_out + tile_base, values); + } + } + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + + // Dequeue and scan tiles of items as part of a dynamic chained scan + // with Init functor + _CCCL_DEVICE _CCCL_FORCEINLINE AgentScanByKey( + TempStorage& storage, + KeysInputIteratorT d_keys_in, + KeyT* d_keys_prev_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + EqualityOp equality_op, + ScanOpT scan_op, + InitValueT init_value) + : storage(storage.Alias()) + , d_keys_in(d_keys_in) + , d_keys_prev_in(d_keys_prev_in) + , d_values_in(d_values_in) + , d_values_out(d_values_out) + , inequality_op(equality_op) + , scan_op(scan_op) + , pair_scan_op(scan_op) + , init_value(init_value) + {} + + /** + * Scan tiles of items as part of a dynamic chained scan + * + * @param num_items + * Total number of input items + * + * @param tile_state + * Global tile state descriptor + * + * start_tile + * The starting tile for the current grid + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeRange(OffsetT num_items, ScanTileStateT& tile_state, int start_tile) + { + int tile_idx = static_cast(blockIdx.x); + OffsetT tile_base = OffsetT(ITEMS_PER_TILE) * tile_idx; + OffsetT num_remaining = num_items - tile_base; + + if (num_remaining > ITEMS_PER_TILE) + { + // Not the last tile (full) + ConsumeTile(num_items, num_remaining, tile_idx, tile_base, tile_state); + } + else if (num_remaining > 0) + { + // The last tile (possibly partially-full) + ConsumeTile(num_items, num_remaining, tile_idx, tile_base, tile_state); + } + } +}; +} // namespace detail::scan_by_key + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_segmented_radix_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_segmented_radix_sort.cuh new file mode 100644 index 00000000..b8ec32ca --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_segmented_radix_sort.cuh @@ -0,0 +1,263 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::radix_sort +{ +/** + * This agent will be implementing the `DeviceSegmentedRadixSort` when the + * https://github.com/NVIDIA/cub/issues/383 is addressed. + * + * @tparam IsDescending + * Whether or not the sorted-order is high-to-low + * + * @tparam SegmentedPolicyT + * Chained tuning policy + * + * @tparam KeyT + * Key type + * + * @tparam ValueT + * Value type + * + * @tparam OffsetT + * Signed integer type for global offsets + */ +template +struct AgentSegmentedRadixSort +{ + OffsetT num_items; + + static constexpr int ITEMS_PER_THREAD = SegmentedPolicyT::ITEMS_PER_THREAD; + static constexpr int BLOCK_THREADS = SegmentedPolicyT::BLOCK_THREADS; + static constexpr int RADIX_BITS = SegmentedPolicyT::RADIX_BITS; + static constexpr int RADIX_DIGITS = 1 << RADIX_BITS; + static constexpr int KEYS_ONLY = ::cuda::std::is_same_v; + + using traits = radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + + // Huge segment handlers + using BlockUpsweepT = AgentRadixSortUpsweep; + using DigitScanT = BlockScan; + using BlockDownsweepT = AgentRadixSortDownsweep; + + /// Number of bin-starting offsets tracked per thread + static constexpr int BINS_TRACKED_PER_THREAD = BlockDownsweepT::BINS_TRACKED_PER_THREAD; + + // Small segment handlers + using BlockRadixSortT = + BlockRadixSort; + + using BlockKeyLoadT = BlockLoad; + + using BlockValueLoadT = BlockLoad; + + union _TempStorage + { + // Huge segment handlers + typename BlockUpsweepT::TempStorage upsweep; + typename BlockDownsweepT::TempStorage downsweep; + + struct UnboundBlockSort + { + OffsetT reverse_counts_in[RADIX_DIGITS]; + OffsetT reverse_counts_out[RADIX_DIGITS]; + typename DigitScanT::TempStorage scan; + } unbound_sort; + + // Small segment handlers + typename BlockKeyLoadT::TempStorage keys_load; + typename BlockValueLoadT::TempStorage values_load; + typename BlockRadixSortT::TempStorage sort; + }; + + using TempStorage = Uninitialized<_TempStorage>; + _TempStorage& temp_storage; + + DecomposerT decomposer; + + _CCCL_DEVICE _CCCL_FORCEINLINE + AgentSegmentedRadixSort(OffsetT num_items, TempStorage& temp_storage, DecomposerT decomposer = {}) + : num_items(num_items) + , temp_storage(temp_storage.Alias()) + , decomposer(decomposer) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE void ProcessSinglePass( + int begin_bit, int end_bit, const KeyT* d_keys_in, const ValueT* d_values_in, KeyT* d_keys_out, ValueT* d_values_out) + { + KeyT thread_keys[ITEMS_PER_THREAD]; + ValueT thread_values[ITEMS_PER_THREAD]; + + // For FP64 the difference is: + // Lowest() -> -1.79769e+308 = 00...00b -> TwiddleIn -> -0 = 10...00b + // LOWEST -> -nan = 11...11b -> TwiddleIn -> 0 = 00...00b + + bit_ordered_type default_key_bits = + IsDescending ? traits::min_raw_binary_key(decomposer) : traits::max_raw_binary_key(decomposer); + KeyT oob_default = reinterpret_cast(default_key_bits); + + if (!KEYS_ONLY) + { + BlockValueLoadT(temp_storage.values_load).Load(d_values_in, thread_values, num_items); + + __syncthreads(); + } + + { + BlockKeyLoadT(temp_storage.keys_load).Load(d_keys_in, thread_keys, num_items, oob_default); + + __syncthreads(); + } + + BlockRadixSortT(temp_storage.sort) + .SortBlockedToStriped( + thread_keys, + thread_values, + begin_bit, + end_bit, + bool_constant_v, + bool_constant_v, + decomposer); + + cub::StoreDirectStriped(threadIdx.x, d_keys_out, thread_keys, num_items); + + if (!KEYS_ONLY) + { + cub::StoreDirectStriped(threadIdx.x, d_values_out, thread_values, num_items); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ProcessIterative( + int current_bit, + int pass_bits, + const KeyT* d_keys_in, + const ValueT* d_values_in, + KeyT* d_keys_out, + ValueT* d_values_out) + { + // Upsweep + BlockUpsweepT upsweep(temp_storage.upsweep, d_keys_in, current_bit, pass_bits, decomposer); + upsweep.ProcessRegion(OffsetT{}, num_items); + + __syncthreads(); + + // The count of each digit value in this pass (valid in the first RADIX_DIGITS threads) + OffsetT bin_count[BINS_TRACKED_PER_THREAD]; + upsweep.ExtractCounts(bin_count); + + __syncthreads(); + + if (IsDescending) + { + // Reverse bin counts + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + temp_storage.unbound_sort.reverse_counts_in[bin_idx] = bin_count[track]; + } + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + bin_count[track] = temp_storage.unbound_sort.reverse_counts_in[RADIX_DIGITS - bin_idx - 1]; + } + } + } + + // Scan + // The global scatter base offset for each digit value in this pass + // (valid in the first RADIX_DIGITS threads) + OffsetT bin_offset[BINS_TRACKED_PER_THREAD]; + DigitScanT(temp_storage.unbound_sort.scan).ExclusiveSum(bin_count, bin_offset); + + if (IsDescending) + { + // Reverse bin offsets + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + temp_storage.unbound_sort.reverse_counts_out[threadIdx.x] = bin_offset[track]; + } + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (threadIdx.x * BINS_TRACKED_PER_THREAD) + track; + + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + bin_offset[track] = temp_storage.unbound_sort.reverse_counts_out[RADIX_DIGITS - bin_idx - 1]; + } + } + } + + __syncthreads(); + + // Downsweep + BlockDownsweepT downsweep( + temp_storage.downsweep, + bin_offset, + num_items, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + current_bit, + pass_bits, + decomposer); + downsweep.ProcessRegion(OffsetT{}, num_items); + } +}; +} // namespace detail::radix_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_select_if.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_select_if.cuh new file mode 100644 index 00000000..44558013 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_select_if.cuh @@ -0,0 +1,1074 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * cub::AgentSelectIf implements a stateful abstraction of CUDA thread blocks for participating in device-wide select. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +namespace detail +{ +// TODO(bgruber): remove this when C++20 is the minimum, since then we can pass policy values as NTTP +template > +struct agent_select_if_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int ITEMS_PER_THREAD = ItemsPerThread; + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + + struct detail + { + using delay_constructor_t = DelayConstructorT; + }; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentSelectIfPolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceSelect/DevicePartition") = detail:: + agent_select_if_policy; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::select +{ +template +struct guarded_inequality_op +{ + EqualityOpT op; + int num_remaining; + + template , int> = 0> + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator()(const T& a, const T& b, int idx) noexcept( + ::cuda::std::__is_nothrow_callable_v) + { + if (idx < num_remaining) + { + return !op(a, b); // In bounds + } + + // Flag out-of-bounds items as selected (as they are discounted for in the agent implementation) + return true; + } + + template , int> = 0> + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE bool operator()(const T& a, const T& b, int idx) const + noexcept(::cuda::std::__is_nothrow_callable_v) + { + if (idx < num_remaining) + { + return !op(a, b); // In bounds + } + + // Flag out-of-bounds items as selected (as they are discounted for in the agent implementation) + return true; + } +}; + +template +struct partition_distinct_output_t +{ + using selected_iterator_t = SelectedOutputItT; + using rejected_iterator_t = RejectedOutputItT; + + selected_iterator_t selected_it; + rejected_iterator_t rejected_it; +}; + +template +struct is_partition_distinct_output_t : ::cuda::std::false_type +{}; + +template +struct is_partition_distinct_output_t> + : ::cuda::std::true_type +{}; + +/** + * @brief AgentSelectIf implements a stateful abstraction of CUDA thread blocks for participating in + * device-wide selection + * + * Performs functor-based selection if SelectOpT functor type != NullType + * Otherwise performs flag-based selection if FlagsInputIterator's value type != NullType + * Otherwise performs discontinuity selection (keep unique) + * + * @tparam AgentSelectIfPolicyT + * Parameterized AgentSelectIfPolicy tuning policy type + * + * @tparam InputIteratorT + * Random-access input iterator type for selection items + * + * @tparam FlagsInputIteratorT + * Random-access input iterator type for selections (NullType* if a selection functor or + * discontinuity flagging is to be used for selection) + * + * @tparam OutputIteratorWrapperT + * Either a random-access iterator or an instance of the `partition_distinct_output_t` template. + * + * @tparam SelectOpT + * Selection operator type (NullType if selections or discontinuity flagging is to be used for + * selection) + * + * @tparam EqualityOpT + * Equality operator type (NullType if selection functor or selections is to be used for + * selection) + * + * @tparam OffsetT + * Signed integer type for offsets within a partition + * + * @tparam StreamingContextT + * Type providing the context information for the current partition, with the following member functions: + * input_offset() -> base offset for the input (and flags) iterator + * is_first_partition() -> [Select::Unique-only] whether this is the first partition + * num_previously_selected() -> base offset for the output iterator for selected items + * num_previously_rejected() -> base offset for the output iterator for rejected items (partition only) + * num_total_items() -> total number of items across all partitions (partition only) + * update_num_selected(d_num_sel_out, num_selected) -> invoked by last CTA with number of selected + * + * @tparam SelectImpl SelectionOpt + * SelectImpl indicating whether to partition, just selection or selection where the memory for the input and + * output may alias each other. + */ +template +struct AgentSelectIf +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + using ScanTileStateT = ScanTileState; + + // Indicates whether the BlockLoad algorithm uses shared memory to load or exchange the data + static constexpr bool loads_via_smem = + !(AgentSelectIfPolicyT::LOAD_ALGORITHM == BLOCK_LOAD_DIRECT + || AgentSelectIfPolicyT::LOAD_ALGORITHM == BLOCK_LOAD_STRIPED + || AgentSelectIfPolicyT::LOAD_ALGORITHM == BLOCK_LOAD_VECTORIZE); + + // If this may be an *in-place* stream compaction, we need to ensure that all of a tile's items have been loaded + // before signalling a subsequent thread block's partial or inclusive state, hence we need a store release when + // updating a tile state. Similarly, we need to make sure that the load of previous tile states precede writing of + // the stream-compacted items and, hence, we need a load acquire when reading those tile states. + static constexpr MemoryOrder memory_order = + ((SelectionOpt == SelectImpl::SelectPotentiallyInPlace) && (!loads_via_smem)) + ? MemoryOrder::acquire_release + : MemoryOrder::relaxed; + + // If we need to enforce memory order for in-place stream compaction, wrap the default decoupled look-back tile + // state in a helper class that enforces memory order on reads and writes + using MemoryOrderedTileStateT = tile_state_with_memory_order; + + // The input value type + using InputT = it_value_t; + + // The flag value type + using FlagT = it_value_t; + + // Constants + enum + { + USE_SELECT_OP, + USE_SELECT_FLAGS, + USE_DISCONTINUITY, + USE_STENCIL_WITH_OP + }; + + static constexpr ::cuda::std::int32_t BLOCK_THREADS = AgentSelectIfPolicyT::BLOCK_THREADS; + static constexpr ::cuda::std::int32_t ITEMS_PER_THREAD = AgentSelectIfPolicyT::ITEMS_PER_THREAD; + static constexpr ::cuda::std::int32_t TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD; + static constexpr bool TWO_PHASE_SCATTER = (ITEMS_PER_THREAD > 1); + + static constexpr bool has_select_op = (!::cuda::std::is_same_v); + static constexpr bool has_flags_it = (!::cuda::std::is_same_v); + static constexpr bool use_stencil_with_op = has_select_op && has_flags_it; + static constexpr auto SELECT_METHOD = + use_stencil_with_op ? USE_STENCIL_WITH_OP + : has_select_op ? USE_SELECT_OP + : has_flags_it ? USE_SELECT_FLAGS + : USE_DISCONTINUITY; + + // Cache-modified Input iterator wrapper type (for applying cache modifier) for items + // Wrap the native input pointer with CacheModifiedValuesInputIterator + // or directly use the supplied input iterator type + using WrappedInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + InputIteratorT>; + + // Cache-modified Input iterator wrapper type (for applying cache modifier) for values + // Wrap the native input pointer with CacheModifiedValuesInputIterator + // or directly use the supplied input iterator type + using WrappedFlagsInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + CacheModifiedInputIterator, + FlagsInputIteratorT>; + + // Parameterized BlockLoad type for input data + using BlockLoadT = BlockLoad; + + // Parameterized BlockLoad type for flags + using BlockLoadFlags = BlockLoad; + + // Parameterized BlockDiscontinuity type for items + using BlockDiscontinuityT = BlockDiscontinuity; + + // Parameterized BlockScan type + using BlockScanT = BlockScan; + + // Callback type for obtaining tile prefix during block scan + using DelayConstructorT = typename AgentSelectIfPolicyT::detail::delay_constructor_t; + using TilePrefixCallbackOpT = + TilePrefixCallbackOp, MemoryOrderedTileStateT, DelayConstructorT>; + + // Item exchange type + using ItemExchangeT = InputT[TILE_ITEMS]; + + // Shared memory type for this thread block + union _TempStorage + { + struct ScanStorage + { + // Smem needed for tile scanning + typename BlockScanT::TempStorage scan; + + // Smem needed for cooperative prefix callback + typename TilePrefixCallbackOpT::TempStorage prefix; + + // Smem needed for discontinuity detection + typename BlockDiscontinuityT::TempStorage discontinuity; + } scan_storage; + + // Smem needed for loading items + typename BlockLoadT::TempStorage load_items; + + // Smem needed for loading values + typename BlockLoadFlags::TempStorage load_flags; + + // Smem needed for compacting items (allows non POD items in this union) + Uninitialized raw_exchange; + }; + + // Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + _TempStorage& temp_storage; ///< Reference to temp_storage + WrappedInputIteratorT d_in; ///< Input items + OutputIteratorWrapperT d_selected_out; ///< Output iterator for the selected items + WrappedFlagsInputIteratorT d_flags_in; ///< Input selection flags (if applicable) + EqualityOpT equality_op; ///< T equality operator + SelectOpT select_op; ///< Selection operator + OffsetT num_items; ///< Total number of input items + + // Note: This is a const reference because we have seen double-digit percentage perf regressions otherwise + const StreamingContextT& streaming_context; ///< Context for the current partition + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + + /** + * @param temp_storage + * Reference to temp_storage + * + * @param d_in + * Input data + * + * @param d_flags_in + * Input selection flags (if applicable) + * + * @param d_selected_out + * Output data + * + * @param select_op + * Selection operator + * + * @param equality_op + * Equality operator + * + * @param num_items + * Total number of input items + * + * @param streaming_context + * Context for the current partition + */ + _CCCL_DEVICE _CCCL_FORCEINLINE AgentSelectIf( + TempStorage& temp_storage, + InputIteratorT d_in, + FlagsInputIteratorT d_flags_in, + OutputIteratorWrapperT d_selected_out, + SelectOpT select_op, + EqualityOpT equality_op, + OffsetT num_items, + const StreamingContextT& streaming_context) + : temp_storage(temp_storage.Alias()) + , d_in(d_in) + , d_selected_out(d_selected_out) + , d_flags_in(d_flags_in) + , equality_op(equality_op) + , select_op(select_op) + , num_items(num_items) + , streaming_context(streaming_context) + {} + + //--------------------------------------------------------------------- + // Utility methods for initializing the selections + //--------------------------------------------------------------------- + + /** + * Initialize selections (specialized for selection operator) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InitializeSelections( + OffsetT /*tile_offset*/, + OffsetT num_tile_items, + InputT (&items)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + constant_t /*select_method*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + // Out-of-bounds items are selection_flags + selection_flags[ITEM] = 1; + + if (!IS_LAST_TILE || (static_cast(threadIdx.x * ITEMS_PER_THREAD + ITEM) < num_tile_items)) + { + selection_flags[ITEM] = static_cast(select_op(items[ITEM])); + } + } + } + + /** + * Initialize selections (specialized for selection_op applied to d_flags_in) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InitializeSelections( + OffsetT tile_offset, + OffsetT num_tile_items, + InputT (& /*items*/)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + constant_t /*select_method*/) + { + __syncthreads(); + + FlagT flags[ITEMS_PER_THREAD]; + if (IS_LAST_TILE) + { + // Initialize the out-of-bounds flags + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + selection_flags[ITEM] = true; + } + // Guarded loads + BlockLoadFlags(temp_storage.load_flags) + .Load((d_flags_in + streaming_context.input_offset()) + tile_offset, flags, num_tile_items); + } + else + { + BlockLoadFlags(temp_storage.load_flags).Load((d_flags_in + streaming_context.input_offset()) + tile_offset, flags); + } + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + // Set selection_flags for out-of-bounds items + if ((!IS_LAST_TILE) || (static_cast(threadIdx.x * ITEMS_PER_THREAD + ITEM) < num_tile_items)) + { + selection_flags[ITEM] = static_cast(select_op(flags[ITEM])); + } + } + } + + /** + * Initialize selections (specialized for valid flags) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InitializeSelections( + OffsetT tile_offset, + OffsetT num_tile_items, + InputT (& /*items*/)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + constant_t /*select_method*/) + { + __syncthreads(); + + FlagT flags[ITEMS_PER_THREAD]; + + if (IS_LAST_TILE) + { + // Out-of-bounds items are selection_flags + BlockLoadFlags(temp_storage.load_flags) + .Load((d_flags_in + streaming_context.input_offset()) + tile_offset, flags, num_tile_items, 1); + } + else + { + BlockLoadFlags(temp_storage.load_flags).Load((d_flags_in + streaming_context.input_offset()) + tile_offset, flags); + } + + // Convert flag type to selection_flags type + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + selection_flags[ITEM] = static_cast(flags[ITEM]); + } + } + + /** + * Initialize selections (specialized for discontinuity detection) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InitializeSelections( + OffsetT tile_offset, + OffsetT num_tile_items, + InputT (&items)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + constant_t /*select_method*/) + { + // We previously invoked the equality operator on out-of-bounds items + // While fixing that issue there were some performance regressions that we had to work around + // To avoid invoking equality operator on unexpected values we are doing on of two things: + // (1) for primitive types AND ::cuda::std::equal_to: we compare all items, including out-of-bounds items and later + // correct the flags of the out-of-bounds items + // (2) otherwise, we are guarding against invoking the equality operator on out-of-bounds items + static constexpr bool use_flag_fixup_code_path = + ::cuda::std::is_arithmetic_v + && (::cuda::std::is_same_v> + || ::cuda::std::is_same_v>); + + if (IS_FIRST_TILE && streaming_context.is_first_partition()) + { + __syncthreads(); + + if constexpr (IS_LAST_TILE && !use_flag_fixup_code_path) + { + // Use custom flag operator to additionally flag the first out-of-bounds item + guarded_inequality_op flag_op{equality_op, num_tile_items}; + + // Set head selection_flags. First tile sets the first flag for the first item + BlockDiscontinuityT(temp_storage.scan_storage.discontinuity).FlagHeads(selection_flags, items, flag_op); + } + else + { + // Set head selection_flags. First tile sets the first flag for the first item + BlockDiscontinuityT(temp_storage.scan_storage.discontinuity) + .FlagHeads(selection_flags, items, InequalityWrapper{equality_op}); + } + } + else + { + InputT tile_predecessor; + if (threadIdx.x == 0) + { + tile_predecessor = d_in[tile_offset + streaming_context.input_offset() - 1]; + } + + __syncthreads(); + + if constexpr (IS_LAST_TILE && !use_flag_fixup_code_path) + { + // Use custom flag operator to additionally flag the first out-of-bounds item + guarded_inequality_op flag_op{equality_op, num_tile_items}; + + // Set head selection_flags. First tile sets the first flag for the first item + BlockDiscontinuityT(temp_storage.scan_storage.discontinuity) + .FlagHeads(selection_flags, items, flag_op, tile_predecessor); + } + else + { + // Set head selection_flags. First tile sets the first flag for the first item + BlockDiscontinuityT(temp_storage.scan_storage.discontinuity) + .FlagHeads(selection_flags, items, InequalityWrapper{equality_op}, tile_predecessor); + } + } + + // For primitive types with default equality operator, we need to fix up the flags for the out-of-bounds items + if constexpr (use_flag_fixup_code_path) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + // Set selection_flags for out-of-bounds items + if ((IS_LAST_TILE) && (OffsetT(threadIdx.x * ITEMS_PER_THREAD) + ITEM >= num_tile_items)) + { + selection_flags[ITEM] = 1; + } + } + } + } + + //--------------------------------------------------------------------- + // Scatter utility methods + //--------------------------------------------------------------------- + + /** + * Scatter flagged items to output offsets (specialized for direct scattering). + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterSelectedDirect( + InputT (&items)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + OffsetT (&selection_indices)[ITEMS_PER_THREAD], + OffsetT num_selections) + { + // Scatter flagged items + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + if (selection_flags[ITEM]) + { + // selection_indices could potentially overflow on the last tile if that's close to INT_MAX, + // so in the streaming invocation we split at INT_MAX round down to a integer multiple of TILE_ITEMS. + if ((!IS_LAST_TILE) || selection_indices[ITEM] < num_selections) + { + *((d_selected_out + streaming_context.num_previously_selected()) + selection_indices[ITEM]) = items[ITEM]; + } + } + } + } + + /** + * @brief Scatter flagged items to output offsets (specialized for two-phase scattering) + * + * @param num_tile_items + * Number of valid items in this tile + * + * @param num_tile_selections + * Number of selections in this tile + * + * @param num_selections_prefix + * Total number of selections prior to this tile + * + * @param num_rejected_prefix + * Total number of rejections prior to this tile + * + * @param is_keep_rejects + * Marker type indicating whether to keep rejected items in the second partition + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterSelectedTwoPhase( + InputT (&items)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + OffsetT (&selection_indices)[ITEMS_PER_THREAD], + int num_tile_selections, + OffsetT num_selections_prefix) + { + __syncthreads(); + + // Compact and scatter items + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + int local_scatter_offset = selection_indices[ITEM] - num_selections_prefix; + if (selection_flags[ITEM]) + { + temp_storage.raw_exchange.Alias()[local_scatter_offset] = items[ITEM]; + } + } + + __syncthreads(); + + for (int item = static_cast(threadIdx.x); item < num_tile_selections; item += BLOCK_THREADS) + { + *((d_selected_out + streaming_context.num_previously_selected()) + + (num_selections_prefix + item)) = // NOLINT(bugprone-misplaced-widening-cast) + temp_storage.raw_exchange.Alias()[item]; + } + } + + /** + * @brief Scatter flagged items. Specialized for selection algorithm that simply discards rejected items + * + * @param num_tile_items + * Number of valid items in this tile + * + * @param num_tile_selections + * Number of selections in this tile + * + * @param num_selections_prefix + * Total number of selections prior to this tile + * + * @param num_rejected_prefix + * Total number of rejections prior to this tile + * + * @param num_selections + * Total number of selections including this tile + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Scatter( + InputT (&items)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + OffsetT (&selection_indices)[ITEMS_PER_THREAD], + int num_tile_items, + int num_tile_selections, + OffsetT num_selections_prefix, + OffsetT num_rejected_prefix, + OffsetT num_selections, + ::cuda::std::false_type /*is_keep_rejects*/) + { + // Do a two-phase scatter if two-phase is enabled and the average number of selection_flags items per thread is + // greater than one + if (TWO_PHASE_SCATTER && (num_tile_selections > BLOCK_THREADS)) + { + ScatterSelectedTwoPhase( + items, selection_flags, selection_indices, num_tile_selections, num_selections_prefix); + } + else + { + ScatterSelectedDirect(items, selection_flags, selection_indices, num_selections); + } + } + + /** + * @brief Scatter flagged items. Specialized for partitioning algorithm that writes rejected items to a second + * partition. + * + * @param num_tile_items + * Number of valid items in this tile + * + * @param num_tile_selections + * Number of selections in this tile + * + * @param num_selections_prefix + * Total number of selections prior to this tile + * + * @param num_rejected_prefix + * Total number of rejections prior to this tile + * + * @param is_keep_rejects + * Marker type indicating whether to keep rejected items in the second partition + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Scatter( + InputT (&items)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + OffsetT (&selection_indices)[ITEMS_PER_THREAD], + int num_tile_items, + int num_tile_selections, + OffsetT num_selections_prefix, + OffsetT num_rejected_prefix, + OffsetT num_selections, + ::cuda::std::true_type /*is_keep_rejects*/) + { + __syncthreads(); + + int tile_num_rejections = num_tile_items - num_tile_selections; + + // Scatter items to shared memory (rejections first) + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + int item_idx = (threadIdx.x * ITEMS_PER_THREAD) + ITEM; + int local_selection_idx = selection_indices[ITEM] - num_selections_prefix; + int local_rejection_idx = item_idx - local_selection_idx; + int local_scatter_offset = + (selection_flags[ITEM]) ? tile_num_rejections + local_selection_idx : local_rejection_idx; + + temp_storage.raw_exchange.Alias()[local_scatter_offset] = items[ITEM]; + } + + // Ensure all threads finished scattering to shared memory + __syncthreads(); + + // Gather items from shared memory and scatter to global + ScatterPartitionsToGlobal( + num_tile_items, tile_num_rejections, num_selections_prefix, num_rejected_prefix, d_selected_out); + } + + /** + * @brief Second phase of scattering partitioned items to global memory. Specialized for partitioning to two + * distinct partitions. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterPartitionsToGlobal( + int num_tile_items, + int tile_num_rejections, + OffsetT num_selections_prefix, + OffsetT num_rejected_prefix, + partition_distinct_output_t partitioned_out_wrapper) + { + auto selected_out_it = partitioned_out_wrapper.selected_it + streaming_context.num_previously_selected(); + auto rejected_out_it = partitioned_out_wrapper.rejected_it + streaming_context.num_previously_rejected(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + int item_idx = (ITEM * BLOCK_THREADS) + threadIdx.x; + int rejection_idx = item_idx; + int selection_idx = item_idx - tile_num_rejections; + OffsetT scatter_offset = + (item_idx < tile_num_rejections) ? num_rejected_prefix + rejection_idx : num_selections_prefix + selection_idx; + + InputT item = temp_storage.raw_exchange.Alias()[item_idx]; + + if (!IS_LAST_TILE || (item_idx < num_tile_items)) + { + if (item_idx >= tile_num_rejections) + { + selected_out_it[scatter_offset] = item; + } + else + { + rejected_out_it[scatter_offset] = item; + } + } + } + } + + /** + * @brief Second phase of scattering partitioned items to global memory. Specialized for partitioning to a single + * iterator, where selected items are written in order from the beginning of the iterator and rejected items are + * writtem from the iterators end backwards. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterPartitionsToGlobal( + int num_tile_items, + int tile_num_rejections, + OffsetT num_selections_prefix, + OffsetT num_rejected_prefix, + PartitionedOutputItT partitioned_out_it) + { + using total_offset_t = typename StreamingContextT::total_num_items_t; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + int item_idx = (ITEM * BLOCK_THREADS) + threadIdx.x; + int rejection_idx = item_idx; + int selection_idx = item_idx - tile_num_rejections; + total_offset_t scatter_offset = + (item_idx < tile_num_rejections) + ? (streaming_context.num_total_items(num_items) - streaming_context.num_previously_rejected() + - static_cast(num_rejected_prefix) - static_cast(rejection_idx) + - total_offset_t{1}) + : (streaming_context.num_previously_selected() + static_cast(num_selections_prefix) + + static_cast(selection_idx)); + + InputT item = temp_storage.raw_exchange.Alias()[item_idx]; + if (!IS_LAST_TILE || (item_idx < num_tile_items)) + { + partitioned_out_it[scatter_offset] = item; + } + } + } + + //--------------------------------------------------------------------- + // Cooperatively scan a device-wide sequence of tiles with other CTAs + //--------------------------------------------------------------------- + + /** + * @brief Process first tile of input (dynamic chained scan). + * + * @param num_tile_items + * Number of input items comprising this tile + * + * @param tile_offset + * Tile offset + * + * @param tile_state_wrapper + * A global tile state descriptor wrapped in a MemoryOrderedTileStateT that ensures consistent memory order across + * all tile status updates and loads + * + * @return The running count of selections (including this tile) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE OffsetT + ConsumeFirstTile(int num_tile_items, OffsetT tile_offset, MemoryOrderedTileStateT& tile_state_wrapper) + { + InputT items[ITEMS_PER_THREAD]; + OffsetT selection_flags[ITEMS_PER_THREAD]; + OffsetT selection_indices[ITEMS_PER_THREAD]; + + // Load items + if (IS_LAST_TILE) + { + BlockLoadT(temp_storage.load_items) + .Load((d_in + streaming_context.input_offset()) + tile_offset, items, num_tile_items); + } + else + { + BlockLoadT(temp_storage.load_items).Load((d_in + streaming_context.input_offset()) + tile_offset, items); + } + + // Initialize selection_flags + InitializeSelections( + tile_offset, num_tile_items, items, selection_flags, constant_v); + + // Ensure temporary storage used during block load can be reused + // Also, in case of in-place stream compaction, this is needed to order the loads of + // *all threads of this thread block* before the st.release of the thread writing this thread block's tile state + __syncthreads(); + + // Exclusive scan of selection_flags + OffsetT num_tile_selections; + BlockScanT(temp_storage.scan_storage.scan).ExclusiveSum(selection_flags, selection_indices, num_tile_selections); + + if (threadIdx.x == 0) + { + // Update tile status if this is not the last tile + if (!IS_LAST_TILE) + { + tile_state_wrapper.SetInclusive(0, num_tile_selections); + } + } + + // Discount any out-of-bounds selections + if (IS_LAST_TILE) + { + num_tile_selections -= (TILE_ITEMS - num_tile_items); + } + + // Scatter flagged items + Scatter( + items, + selection_flags, + selection_indices, + num_tile_items, + num_tile_selections, + 0, + 0, + num_tile_selections, + bool_constant_v < SelectionOpt == SelectImpl::Partition >); + + return num_tile_selections; + } + + /** + * @brief Process subsequent tile of input (dynamic chained scan). + * + * @param num_tile_items + * Number of input items comprising this tile + * + * @param tile_idx + * Tile index + * + * @param tile_offset + * Tile offset + * + * @param tile_state_wrapper + * A global tile state descriptor wrapped in a MemoryOrderedTileStateT that ensures consistent memory order across + * all tile status updates and loads + * + * @return The running count of selections (including this tile) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE OffsetT ConsumeSubsequentTile( + int num_tile_items, int tile_idx, OffsetT tile_offset, MemoryOrderedTileStateT& tile_state_wrapper) + { + InputT items[ITEMS_PER_THREAD]; + OffsetT selection_flags[ITEMS_PER_THREAD]; + OffsetT selection_indices[ITEMS_PER_THREAD]; + + // Load items + if (IS_LAST_TILE) + { + BlockLoadT(temp_storage.load_items) + .Load((d_in + streaming_context.input_offset()) + tile_offset, items, num_tile_items); + } + else + { + BlockLoadT(temp_storage.load_items).Load((d_in + streaming_context.input_offset()) + tile_offset, items); + } + + // Initialize selection_flags + InitializeSelections( + tile_offset, num_tile_items, items, selection_flags, constant_v); + + // Ensure temporary storage used during block load can be reused + // Also, in case of in-place stream compaction, this is needed to order the loads of + // *all threads of this thread block* before the st.release of the thread writing this thread block's tile state + __syncthreads(); + + // Exclusive scan of values and selection_flags + TilePrefixCallbackOpT prefix_op( + tile_state_wrapper, temp_storage.scan_storage.prefix, ::cuda::std::plus<>{}, tile_idx); + BlockScanT(temp_storage.scan_storage.scan).ExclusiveSum(selection_flags, selection_indices, prefix_op); + + OffsetT num_tile_selections = prefix_op.GetBlockAggregate(); + OffsetT num_selections = prefix_op.GetInclusivePrefix(); + OffsetT num_selections_prefix = prefix_op.GetExclusivePrefix(); + OffsetT num_rejected_prefix = tile_offset - num_selections_prefix; + + // Discount any out-of-bounds selections + if (IS_LAST_TILE) + { + int num_discount = TILE_ITEMS - num_tile_items; + num_selections -= num_discount; + num_tile_selections -= num_discount; + } + + // note (only applies to in-place stream compaction): We can avoid having to introduce explicit memory order between + // the look-back (i.e., loading previous tiles' states) and scattering items (which means, potentially overwriting + // previous tiles' input items, in case of in-place compaction), because this is implicitly ensured through + // execution dependency: The scatter stage requires the offset from the prefix-sum and it can only know the + // prefix-sum after having read that from the decoupled look-back. Scatter flagged items + Scatter( + items, + selection_flags, + selection_indices, + num_tile_items, + num_tile_selections, + num_selections_prefix, + num_rejected_prefix, + num_selections, + bool_constant_v < SelectionOpt == SelectImpl::Partition >); + + return num_selections; + } + + /** + * @brief Process a tile of input + * + * @param num_tile_items + * Number of input items comprising this tile + * + * @param tile_idx + * Tile index + * + * @param tile_offset + * Tile offset + * + * @param tile_state_wrapper + * A global tile state descriptor wrapped in a MemoryOrderedTileStateT that ensures consistent memory order across + * all tile status updates and loads + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE OffsetT + ConsumeTile(int num_tile_items, int tile_idx, OffsetT tile_offset, MemoryOrderedTileStateT& tile_state_wrapper) + { + OffsetT num_selections; + if (tile_idx == 0) + { + num_selections = ConsumeFirstTile(num_tile_items, tile_offset, tile_state_wrapper); + } + else + { + num_selections = ConsumeSubsequentTile(num_tile_items, tile_idx, tile_offset, tile_state_wrapper); + } + + return num_selections; + } + + /** + * @brief Scan tiles of items as part of a dynamic chained scan + * + * @param num_tiles + * Total number of input tiles + * + * @param tile_state + * Global tile state descriptor + * + * @param d_num_selected_out + * Output total number selection_flags + * + * @tparam NumSelectedIteratorT + * Output iterator type for recording number of items selection_flags + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeRange(int num_tiles, ScanTileStateT& tile_state, NumSelectedIteratorT d_num_selected_out) + { + // Ensure consistent memory order across all tile status updates and loads + auto tile_state_wrapper = MemoryOrderedTileStateT{tile_state}; + + // Blocks are launched in increasing order, so just assign one tile per block + // TODO (elstehle): replacing this term with just `blockIdx.x` degrades perf for partition. Once we get to re-tune + // the algorithm, we want to replace this term with `blockIdx.x` + int tile_idx{}; + if constexpr (SELECT_METHOD != USE_DISCONTINUITY) + { + tile_idx = static_cast((blockIdx.x * gridDim.y) + blockIdx.y); // Current tile index + } + else + { + tile_idx = static_cast(blockIdx.x); // Current tile index + } + OffsetT tile_offset = static_cast(tile_idx) * static_cast(TILE_ITEMS); + + if (tile_idx < num_tiles - 1) + { + // Not the last tile (full) + ConsumeTile(TILE_ITEMS, tile_idx, tile_offset, tile_state_wrapper); + } + else + { + // The last tile (possibly partially-full) + OffsetT num_remaining = num_items - tile_offset; + OffsetT num_selections = ConsumeTile(num_remaining, tile_idx, tile_offset, tile_state_wrapper); + + if (threadIdx.x == 0) + { + // Update the number of selected items with this partition's selections + streaming_context.update_num_selected(d_num_selected_out, num_selections); + } + } + } +}; +} // namespace detail::select + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_sub_warp_merge_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_sub_warp_merge_sort.cuh new file mode 100644 index 00000000..5346d959 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_sub_warp_merge_sort.cuh @@ -0,0 +1,349 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +// TODO(bgruber): drop in CCCL 4.0 +template +struct agent_sub_warp_merge_sort_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int WARP_THREADS = WarpThreadsArg; + static constexpr int ITEMS_PER_THREAD = ItemsPerThreadArg; + static constexpr int ITEMS_PER_TILE = WARP_THREADS * ITEMS_PER_THREAD; + static constexpr int SEGMENTS_PER_BLOCK = BLOCK_THREADS / WARP_THREADS; + + static constexpr cub::WarpLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithmArg; + static constexpr cub::CacheLoadModifier LOAD_MODIFIER = LoadModifierArg; + static constexpr cub::WarpStoreAlgorithm STORE_ALGORITHM = StoreAlgorithmArg; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template +using AgentSubWarpMergeSortPolicy + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceSegmentedSort") = detail::agent_sub_warp_merge_sort_policy< + ThreadsPerBlock, + WarpThreadsArg, + ItemsPerThreadArg, + LoadAlgorithmArg, + LoadModifierArg, + StoreAlgorithmArg>; + +namespace detail::sub_warp_merge_sort +{ +/** + * @brief AgentSubWarpSort implements a sub-warp merge sort. + * + * This agent can work with any power of two number of threads, not exceeding + * 32. The number of threads is defined in the `PolicyT::WARP_THREADS`. Virtual + * warp of `PolicyT::WARP_THREADS` will efficiently load data using + * `PolicyT::LOAD_ALGORITHM`, sort it using `WarpMergeSort`, and store it back + * using `PolicyT::STORE_ALGORITHM`. + * + * @tparam IS_DESCENDING + * Whether or not the sorted-order is high-to-low + * + * @tparam PolicyT + * Chained tuning policy + * + * @tparam KeyT + * Key type + * + * @tparam ValueT + * Value type + * + * @tparam OffsetT + * Signed integer type for global offsets + */ +template +class AgentSubWarpSort +{ + using traits = detail::radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + + struct BinaryOpT + { + template + _CCCL_DEVICE bool operator()(T lhs, T rhs) const noexcept + { + if constexpr (IS_DESCENDING) + { + return lhs > rhs; + } + else + { + return lhs < rhs; + } + _CCCL_UNREACHABLE(); + } + +#if _CCCL_HAS_NVFP16() + _CCCL_DEVICE bool operator()(__half lhs, __half rhs) const noexcept + { + // Need to explicitly cast to float for SM <= 52. + if constexpr (IS_DESCENDING) + { + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_53, (return __hgt(lhs, rhs);), (return __half2float(lhs) > __half2float(rhs);)); + } + else + { + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_53, (return __hlt(lhs, rhs);), (return __half2float(lhs) < __half2float(rhs);)); + } + _CCCL_UNREACHABLE(); + } +#endif // _CCCL_HAS_NVFP16() + +#if _CCCL_HAS_NVBF16() + _CCCL_DEVICE bool operator()(__nv_bfloat16 lhs, __nv_bfloat16 rhs) const noexcept + { + // Need to explicitly cast to float for SM < 80. + if constexpr (IS_DESCENDING) + { + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_80, (return __hgt(lhs, rhs);), (return __bfloat162float(lhs) > __bfloat162float(rhs);)); + } + else + { + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_80, (return __hlt(lhs, rhs);), (return __bfloat162float(lhs) < __bfloat162float(rhs);)); + } + _CCCL_UNREACHABLE(); + } +#endif // _CCCL_HAS_NVBF16() + }; + +#if _CCCL_HAS_NVFP16() + _CCCL_DEVICE static bool equal(__half lhs, __half rhs) + { + // Need to explicitly cast to float for SM <= 52. + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_53, (return __heq(lhs, rhs);), (return __half2float(lhs) == __half2float(rhs);)); + } +#endif // _CCCL_HAS_NVFP16() + +#if _CCCL_HAS_NVBF16() + _CCCL_DEVICE static bool equal(__nv_bfloat16 lhs, __nv_bfloat16 rhs) + { + // Need to explicitly cast to float for SM < 80. + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_80, (return __heq(lhs, rhs);), (return __bfloat162float(lhs) == __bfloat162float(rhs);)); + } +#endif // _CCCL_HAS_NVBF16() + + template + _CCCL_DEVICE static bool equal(T lhs, T rhs) + { + return lhs == rhs; + } + +public: + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + + using WarpMergeSortT = WarpMergeSort; + + using KeysLoadItT = try_make_cache_modified_iterator_t; + using ItemsLoadItT = try_make_cache_modified_iterator_t; + + using WarpLoadKeysT = cub::WarpLoad; + using WarpLoadItemsT = + cub::WarpLoad; + + using WarpStoreKeysT = + cub::WarpStore; + using WarpStoreItemsT = + cub::WarpStore; + + union _TempStorage + { + typename WarpLoadKeysT::TempStorage load_keys; + typename WarpLoadItemsT::TempStorage load_items; + typename WarpMergeSortT::TempStorage sort; + typename WarpStoreKeysT::TempStorage store_keys; + typename WarpStoreItemsT::TempStorage store_items; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + _TempStorage& storage; + + _CCCL_DEVICE _CCCL_FORCEINLINE explicit AgentSubWarpSort(TempStorage& temp_storage) + : storage(temp_storage.Alias()) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE void ProcessSegment( + int segment_size, KeysLoadItT keys_input, KeyT* keys_output, ItemsLoadItT values_input, ValueT* values_output) + { + WarpMergeSortT warp_merge_sort(storage.sort); + + if (segment_size < 3) + { + ShortCircuit( + warp_merge_sort.get_linear_tid(), + segment_size, + keys_input, + keys_output, + values_input, + values_output, + BinaryOpT{}); + } + else + { + KeyT keys[PolicyT::ITEMS_PER_THREAD]; + ValueT values[PolicyT::ITEMS_PER_THREAD]; + + KeyT oob_default = [&] { + if constexpr (::cuda::std::is_same_v) + { + // Traits::MAX_KEY for `bool` is 0xFF which is different from `true` and makes + // comparison with oob unreliable. + return !IS_DESCENDING; + } + else + { + // For FP64 the difference is: + // Lowest() -> -1.79769e+308 = 00...00b -> TwiddleIn -> -0 = 10...00b + // LOWEST -> -nan = 11...11b -> TwiddleIn -> 0 = 00...00b + + // Segmented sort doesn't support custom types at the moment. + bit_ordered_type default_key_bits = IS_DESCENDING ? traits::min_raw_binary_key(identity_decomposer_t{}) + : traits::max_raw_binary_key(identity_decomposer_t{}); + return reinterpret_cast(default_key_bits); + } + }(); + + WarpLoadKeysT(storage.load_keys).Load(keys_input, keys, segment_size, oob_default); + __syncwarp(warp_merge_sort.get_member_mask()); + + if (!KEYS_ONLY) + { + WarpLoadItemsT(storage.load_items).Load(values_input, values, segment_size); + + __syncwarp(warp_merge_sort.get_member_mask()); + } + + warp_merge_sort.Sort(keys, values, BinaryOpT{}, segment_size, oob_default); + __syncwarp(warp_merge_sort.get_member_mask()); + + WarpStoreKeysT(storage.store_keys).Store(keys_output, keys, segment_size); + + if (!KEYS_ONLY) + { + __syncwarp(warp_merge_sort.get_member_mask()); + WarpStoreItemsT(storage.store_items).Store(values_output, values, segment_size); + } + } + } + +private: + /** + * This method implements a shortcut for sorting less than three items. + * Only the first thread of a virtual warp is used for soring. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ShortCircuit( + unsigned int linear_tid, + OffsetT segment_size, + KeysLoadItT keys_input, + KeyT* keys_output, + ItemsLoadItT values_input, + ValueT* values_output, + CompareOpT binary_op) + { + if (segment_size == 1) + { + if (linear_tid == 0) + { + if (keys_input.ptr != keys_output) + { + keys_output[0] = keys_input[0]; + } + + if (!KEYS_ONLY) + { + if (values_input.ptr != values_output) + { + values_output[0] = values_input[0]; + } + } + } + } + else if (segment_size == 2) + { + if (linear_tid == 0) + { + KeyT lhs = keys_input[0]; + KeyT rhs = keys_input[1]; + + if (equal(lhs, rhs) || binary_op(lhs, rhs)) + { + keys_output[0] = lhs; + keys_output[1] = rhs; + + if (!KEYS_ONLY) + { + if (values_output != values_input.ptr) + { + values_output[0] = values_input[0]; + values_output[1] = values_input[1]; + } + } + } + else + { + keys_output[0] = rhs; + keys_output[1] = lhs; + + if (!KEYS_ONLY) + { + // values_output might be an alias for values_input, so + // we have to use registers here + + const ValueT lhs_val = values_input[0]; + const ValueT rhs_val = values_input[1]; + + values_output[0] = rhs_val; + values_output[1] = lhs_val; + } + } + } + } + } +}; +} // namespace detail::sub_warp_merge_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_three_way_partition.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_three_way_partition.cuh new file mode 100644 index 00000000..505edc55 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_three_way_partition.cuh @@ -0,0 +1,588 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +namespace detail +{ +// TODO(bgruber): remove this when C++20 is the minimum, since then we can pass policy values as NTTP +template > +struct agent_three_way_partition_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int ITEMS_PER_THREAD = ItemsPerThread; + static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + + struct detail + { + using delay_constructor_t = DelayConstructorT; + }; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentThreeWayPartitionPolicy + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DevicePartition") = detail::agent_three_way_partition_policy< + ThreadsPerBlock, + ItemsPerThread, + LoadAlgorithm, + LoadModifier, + ScanAlgorithm, + DelayConstructorT>; + +namespace detail::three_way_partition +{ +template +struct pair_pack_t +{ + OffsetT x, y; + + _CCCL_DEVICE pair_pack_t operator+(const pair_pack_t& other) const + { + return {x + other.x, y + other.y}; + } +}; + +template +struct accumulator_pack_base_t +{ + using pack_t = pair_pack_t; + + _CCCL_DEVICE static pack_t pack(OffsetT f, OffsetT s) + { + return {f, s}; + } + _CCCL_DEVICE static OffsetT first(pack_t packed) + { + return packed.x; + } + _CCCL_DEVICE static OffsetT second(pack_t packed) + { + return packed.y; + } +}; + +template +struct accumulator_pack_base_t> +{ + using pack_t = uint64_t; + + _CCCL_DEVICE static pack_t pack(OffsetT f, OffsetT s) + { + return (static_cast(f) << 32) | static_cast(s); + } + + _CCCL_DEVICE static OffsetT first(pack_t packed) + { + return static_cast(packed >> 32); + } + + _CCCL_DEVICE static OffsetT second(pack_t packed) + { + return static_cast(packed & 0xFFFFFFFF); + } +}; + +template +struct accumulator_pack_t : accumulator_pack_base_t +{ + using base = accumulator_pack_base_t; + using typename base::pack_t; + + _CCCL_DEVICE static void subtract(pack_t& packed, OffsetT val) + { + packed = base::pack(base::first(packed) - val, base::second(packed) - val); + } + + _CCCL_DEVICE static OffsetT sum(pack_t& packed) + { + return base::first(packed) + base::second(packed); + } + + _CCCL_DEVICE static pack_t zero() + { + return {}; + } +}; + +/** + * \brief Implements a device-wide three-way partitioning + * + * Splits input data into three parts based on the selection functors. If the + * first functor selects an item, the algorithm places it in the first part. + * Otherwise, if the second functor selects an item, the algorithm places it in + * the second part. If both functors don't select an item, the algorithm places + * it into the unselected part. + */ +template +struct AgentThreeWayPartition +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + // The input value type + using InputT = it_value_t; + + using AccumPackHelperT = accumulator_pack_t; + using AccumPackT = typename AccumPackHelperT::pack_t; + + // Tile status descriptor interface type + using ScanTileStateT = cub::ScanTileState; + + // Constants + static constexpr int BLOCK_THREADS = PolicyT::BLOCK_THREADS; + static constexpr int ITEMS_PER_THREAD = PolicyT::ITEMS_PER_THREAD; + static constexpr int TILE_ITEMS = BLOCK_THREADS * ITEMS_PER_THREAD; + + using WrappedInputIteratorT = + ::cuda::std::_If<::cuda::std::is_pointer_v, + cub::CacheModifiedInputIterator, + InputIteratorT>; + + // Parameterized BlockLoad type for input data + using BlockLoadT = cub::BlockLoad; + + // Parameterized BlockScan type + using BlockScanT = cub::BlockScan; + + // Callback type for obtaining tile prefix during block scan + using DelayConstructorT = typename PolicyT::detail::delay_constructor_t; + using TilePrefixCallbackOpT = + cub::TilePrefixCallbackOp, ScanTileStateT, DelayConstructorT>; + + // Item exchange type + using ItemExchangeT = InputT[TILE_ITEMS]; + + // Shared memory type for this thread block + union _TempStorage + { + struct ScanStorage + { + // Smem needed for tile scanning + typename BlockScanT::TempStorage scan; + + // Smem needed for cooperative prefix callback + typename TilePrefixCallbackOpT::TempStorage prefix; + } scan_storage; + + // Smem needed for loading items + typename BlockLoadT::TempStorage load_items; + + // Smem needed for compacting items (allows non POD items in this union) + cub::Uninitialized raw_exchange; + }; + + // Alias wrapper allowing storage to be unioned + struct TempStorage : cub::Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + _TempStorage& temp_storage; ///< Reference to temp_storage + WrappedInputIteratorT d_in; ///< Input items + FirstOutputIteratorT d_first_part_out; + SecondOutputIteratorT d_second_part_out; + UnselectedOutputIteratorT d_unselected_out; + SelectFirstPartOp select_first_part_op; + SelectSecondPartOp select_second_part_op; + OffsetT num_items; ///< Total number of input items + + // Note: This is a const reference because we have seen double-digit percentage perf regressions otherwise + const StreamingContextT& streaming_context; ///< Context for the current partition + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + + // Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE AgentThreeWayPartition( + TempStorage& temp_storage, + InputIteratorT d_in, + FirstOutputIteratorT d_first_part_out, + SecondOutputIteratorT d_second_part_out, + UnselectedOutputIteratorT d_unselected_out, + SelectFirstPartOp select_first_part_op, + SelectSecondPartOp select_second_part_op, + OffsetT num_items, + const StreamingContextT& streaming_context) + : temp_storage(temp_storage.Alias()) + , d_in(d_in) + , d_first_part_out(d_first_part_out) + , d_second_part_out(d_second_part_out) + , d_unselected_out(d_unselected_out) + , select_first_part_op(select_first_part_op) + , select_second_part_op(select_second_part_op) + , num_items(num_items) + , streaming_context(streaming_context) + {} + + //--------------------------------------------------------------------- + // Utility methods for initializing the selections + //--------------------------------------------------------------------- + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Initialize( + OffsetT num_tile_items, InputT (&items)[ITEMS_PER_THREAD], AccumPackT (&items_selection_flags)[ITEMS_PER_THREAD]) + { + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + // Out-of-bounds items are selection_flags + items_selection_flags[ITEM] = AccumPackHelperT::pack(1, 1); + + if (!IS_LAST_TILE || (OffsetT(threadIdx.x * ITEMS_PER_THREAD) + ITEM < num_tile_items)) + { + OffsetT first_item_selected = select_first_part_op(items[ITEM]); + items_selection_flags[ITEM] = + AccumPackHelperT::pack(first_item_selected, first_item_selected ? 0 : select_second_part_op(items[ITEM])); + } + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Scatter( + InputT (&items)[ITEMS_PER_THREAD], + AccumPackT (&items_selection_flags)[ITEMS_PER_THREAD], + AccumPackT (&items_selection_indices)[ITEMS_PER_THREAD], + int num_tile_items, + AccumPackT num_tile_selected, + AccumPackT num_tile_selected_prefix, + OffsetT num_rejected_prefix) + { + __syncthreads(); + + const OffsetT num_first_selections_prefix = AccumPackHelperT::first(num_tile_selected_prefix); + const OffsetT num_second_selections_prefix = AccumPackHelperT::second(num_tile_selected_prefix); + + const int first_item_end = AccumPackHelperT::first(num_tile_selected); + const int second_item_end = first_item_end + AccumPackHelperT::second(num_tile_selected); + + // Scatter items to shared memory (rejections first) + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + int item_idx = (threadIdx.x * ITEMS_PER_THREAD) + ITEM; + + const OffsetT first_items_selection_indices = AccumPackHelperT::first(items_selection_indices[ITEM]); + const OffsetT second_items_selection_indices = AccumPackHelperT::second(items_selection_indices[ITEM]); + + if (!IS_LAST_TILE || (item_idx < num_tile_items)) + { + int local_scatter_offset = 0; + + if (AccumPackHelperT::first(items_selection_flags[ITEM])) + { + local_scatter_offset = first_items_selection_indices - num_first_selections_prefix; + } + else if (AccumPackHelperT::second(items_selection_flags[ITEM])) + { + local_scatter_offset = first_item_end + second_items_selection_indices - num_second_selections_prefix; + } + else + { + // Medium item + int local_selection_idx = (first_items_selection_indices - num_first_selections_prefix) + + (second_items_selection_indices - num_second_selections_prefix); + local_scatter_offset = second_item_end + item_idx - local_selection_idx; + } + + temp_storage.raw_exchange.Alias()[local_scatter_offset] = items[ITEM]; + } + } + + __syncthreads(); + + // Gather items from shared memory and scatter to global + // NOLINTBEGIN(bugprone-misplaced-widening-cast) + auto first_base = + d_first_part_out + (streaming_context.num_previously_selected_first() + num_first_selections_prefix); + auto second_base = + d_second_part_out + (streaming_context.num_previously_selected_second() + num_second_selections_prefix); + auto unselected_base = d_unselected_out + (streaming_context.num_previously_rejected() + num_rejected_prefix); + // NOLINTEND(bugprone-misplaced-widening-cast) + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + int item_idx = (ITEM * BLOCK_THREADS) + threadIdx.x; + + if (!IS_LAST_TILE || (item_idx < num_tile_items)) + { + InputT item = temp_storage.raw_exchange.Alias()[item_idx]; + + if (item_idx < first_item_end) + { + first_base[item_idx] = item; + } + else if (item_idx < second_item_end) + { + second_base[item_idx - first_item_end] = item; + } + else + { + int rejection_idx = item_idx - second_item_end; + unselected_base[rejection_idx] = item; + } + } + } + } + + //--------------------------------------------------------------------- + // Cooperatively scan a device-wide sequence of tiles with other CTAs + //--------------------------------------------------------------------- + + /** + * Process first tile of input (dynamic chained scan). + * Returns the running count of selections (including this tile) + * + * @param num_tile_items Number of input items comprising this tile + * @param tile_offset Tile offset + * @param first_tile_state Global tile state descriptor + * @param second_tile_state Global tile state descriptor + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeFirstTile(int num_tile_items, OffsetT tile_offset, ScanTileStateT& tile_state, AccumPackT& num_items_selected) + { + InputT items[ITEMS_PER_THREAD]; + + AccumPackT items_selection_flags[ITEMS_PER_THREAD]; + AccumPackT items_selection_indices[ITEMS_PER_THREAD]; + + // Load items + if constexpr (IS_LAST_TILE) + { + BlockLoadT(temp_storage.load_items) + .Load(d_in + streaming_context.input_offset() + tile_offset, items, num_tile_items); + } + else + { + BlockLoadT(temp_storage.load_items).Load(d_in + streaming_context.input_offset() + tile_offset, items); + } + + // Initialize selection_flags + Initialize(num_tile_items, items, items_selection_flags); + __syncthreads(); + + // Exclusive scan of selection_flags + BlockScanT(temp_storage.scan_storage.scan) + .ExclusiveSum(items_selection_flags, items_selection_indices, num_items_selected); + + if (threadIdx.x == 0) + { + // Update tile status if this is not the last tile + if (!IS_LAST_TILE) + { + tile_state.SetInclusive(0, num_items_selected); + } + } + + // Discount any out-of-bounds selections + if (IS_LAST_TILE) + { + AccumPackHelperT::subtract(num_items_selected, TILE_ITEMS - num_tile_items); + } + + // Scatter flagged items + Scatter( + items, + items_selection_flags, + items_selection_indices, + num_tile_items, + num_items_selected, + // all the prefixes equal to 0 because it's the first tile + AccumPackHelperT::zero(), + 0); + } + + /** + * Process subsequent tile of input (dynamic chained scan). + * Returns the running count of selections (including this tile) + * + * @param num_tile_items Number of input items comprising this tile + * @param tile_idx Tile index + * @param tile_offset Tile offset + * @param first_tile_state Global tile state descriptor + * @param second_tile_state Global tile state descriptor + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ConsumeSubsequentTile( + int num_tile_items, int tile_idx, OffsetT tile_offset, ScanTileStateT& tile_state, AccumPackT& num_items_selected) + { + InputT items[ITEMS_PER_THREAD]; + + AccumPackT items_selected_flags[ITEMS_PER_THREAD]; + AccumPackT items_selected_indices[ITEMS_PER_THREAD]; + + // Load items + if (IS_LAST_TILE) + { + BlockLoadT(temp_storage.load_items) + .Load(d_in + streaming_context.input_offset() + tile_offset, items, num_tile_items); + } + else + { + BlockLoadT(temp_storage.load_items).Load(d_in + streaming_context.input_offset() + tile_offset, items); + } + + // Initialize selection_flags + Initialize(num_tile_items, items, items_selected_flags); + __syncthreads(); + + // Exclusive scan of values and selection_flags + TilePrefixCallbackOpT prefix_op(tile_state, temp_storage.scan_storage.prefix, ::cuda::std::plus<>{}, tile_idx); + + BlockScanT(temp_storage.scan_storage.scan).ExclusiveSum(items_selected_flags, items_selected_indices, prefix_op); + + num_items_selected = prefix_op.GetInclusivePrefix(); + AccumPackT num_items_in_tile_selected = prefix_op.GetBlockAggregate(); + AccumPackT num_items_selected_prefix = prefix_op.GetExclusivePrefix(); + + __syncthreads(); + + OffsetT num_rejected_prefix = (tile_idx * TILE_ITEMS) - AccumPackHelperT::sum(num_items_selected_prefix); + + // Discount any out-of-bounds selections. There are exactly + // TILE_ITEMS - num_tile_items elements like that because we + // marked them as selected in Initialize method. + if (IS_LAST_TILE) + { + const int num_discount = TILE_ITEMS - num_tile_items; + + AccumPackHelperT::subtract(num_items_selected, num_discount); + AccumPackHelperT::subtract(num_items_in_tile_selected, num_discount); + } + + // Scatter flagged items + Scatter( + items, + items_selected_flags, + items_selected_indices, + num_tile_items, + num_items_in_tile_selected, + num_items_selected_prefix, + num_rejected_prefix); + } + + /** + * Process a tile of input + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeTile(int num_tile_items, int tile_idx, OffsetT tile_offset, ScanTileStateT& tile_state, AccumPackT& accum) + { + if (tile_idx == 0) + { + ConsumeFirstTile(num_tile_items, tile_offset, tile_state, accum); + } + else + { + ConsumeSubsequentTile(num_tile_items, tile_idx, tile_offset, tile_state, accum); + } + } + + /** + * Scan tiles of items as part of a dynamic chained scan + * + * @tparam NumSelectedIteratorT + * Output iterator type for recording number of items selection_flags + * + * @param num_tiles + * Total number of input tiles + * + * @param first_tile_state + * Global tile state descriptor + * + * @param second_tile_state + * Global tile state descriptor + * + * @param d_num_selected_out + * Output total number selection_flags + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeRange(int num_tiles, ScanTileStateT& tile_state, NumSelectedIteratorT d_num_selected_out) + { + // Blocks are launched in increasing order, so just assign one tile per block + // Current tile index + const int tile_idx = static_cast(blockIdx.x); + + // Global offset for the current tile + const OffsetT tile_offset = tile_idx * TILE_ITEMS; + + AccumPackT accum; + + if (tile_idx < num_tiles - 1) + { + // Not the last tile (full) + ConsumeTile(TILE_ITEMS, tile_idx, tile_offset, tile_state, accum); + } + else + { + // The last tile (possibly partially-full) + const OffsetT num_remaining = num_items - tile_offset; + + ConsumeTile(num_remaining, tile_idx, tile_offset, tile_state, accum); + + if (threadIdx.x == 0) + { + // Update the number of selected items with this partition's selections + streaming_context.update_num_selected( + d_num_selected_out, AccumPackHelperT::first(accum), AccumPackHelperT::second(accum), num_items); + } + } + } +}; +} // namespace detail::three_way_partition + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_topk.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_topk.cuh new file mode 100644 index 00000000..522712b8 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_topk.cuh @@ -0,0 +1,835 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! @file +//! cub::AgentTopK implements a stateful abstraction of CUDA thread blocks for participating in device-wide topK. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::topk +{ +//! @brief Parameterizable tuning policy type for agent_topk +//! +//! @tparam ThreadsPerBlock +//! Threads per thread block +//! +//! @tparam ItemsPerThread +//! Items per thread (per tile of input) +//! +//! @tparam BitsPerPass +//! Number of bits processed per pass +//! +//! @tparam LoadAlgorithm +//! The BlockLoad algorithm to use +//! +//! @tparam ScanAlgorithm +//! The BlockScan algorithm to use +//! +template +struct agent_topk_policy +{ + static constexpr int threads_per_block = ThreadsPerBlock; + static constexpr int items_per_thread = ItemsPerThread; + static constexpr int bits_per_pass = BitsPerPass; + static constexpr BlockLoadAlgorithm load_algorithm = LoadAlgorithm; + static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; +}; + +template > +struct key_prefix_storage_t; + +template +struct key_prefix_storage_t +{ + using bits_t = typename Traits::UnsignedBits; + bits_t bits; +}; + +// Calculates the number of passes needed for a type T with BitsPerPass bits processed per pass. +template +[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr int calc_num_passes(int bits_per_pass) +{ + return ::cuda::ceil_div(sizeof(T) * 8, bits_per_pass); +} + +template +[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE int calc_num_passes(const int total_bits) +{ + return ::cuda::ceil_div(total_bits, BitsPerPass); +} + +// Calculates the starting bit for a given pass (bit 0 is the least significant (rightmost) bit). +// We process the input from the most to the least significant bit. This way, we can skip some passes in the end. +template +[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr int calc_start_bit(const int pass) +{ + int start_bit = int{sizeof(T)} * 8 - (pass + 1) * BitsPerPass; + if (start_bit < 0) + { + start_bit = 0; + } + return start_bit; +} + +template +[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE int calc_start_bit(const int total_bits, const int pass) +{ + int start_bit = total_bits - (pass + 1) * BitsPerPass; + if (start_bit < 0) + { + start_bit = 0; + } + return start_bit; +} + +// Bit-vector for accumulating prefix digits via funnel shift. Each pass shifts the existing +// contents left by BitsPerPass and ORs the new bucket at the bottom. Sized to hold all +// decomposed bits of KeyT plus headroom for the shift padding of the last pass. +template +struct key_prefix_storage_t +{ + static constexpr int num_words = ::cuda::ceil_div(sizeof(KeyT) * 8 + 31, 32); + unsigned int words[num_words]; + + // Funnel-shifts the entire bit-vector left by `shift` positions and inserts `value` into the + // vacated low bits. Each word receives carry bits from its lower neighbor (high-to-low order + // so each word reads its neighbor's original value). The final word is filled from `value`. + _CCCL_DEVICE _CCCL_FORCEINLINE void shift_or(int shift, unsigned int value) + { + _CCCL_ASSERT(shift > 0 && shift < 32, "shift_or requires 0 < shift < 32"); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = num_words - 1; i > 0; --i) + { + words[i] = __funnelshift_l(words[i - 1], words[i], shift); + } + words[0] = (words[0] << shift) | value; + } +}; + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +set_kth_key_bits(key_prefix_storage_t& prefix, const int pass, const int bin_index) +{ + if constexpr (detail::radix::can_twiddle) + { + using bits_t = typename Traits::UnsignedBits; + const int start_bit = calc_start_bit(pass); + bits_t bucket = bin_index; + prefix.bits |= static_cast(bucket) << start_bit; + } + else + { + prefix.shift_or(BitsPerPass, bin_index); + } +} + +template +struct alignas(128) Counter +{ + // We are processing the items in multiple passes, from most-significant to least-significant bits. In each pass, we + // keep the length of input (`len`) and the `k` of current pass, and update them at the end of the pass. + OutOffsetT k; + OffsetT len; + + // `previous_len` is the length of the input in the previous pass. Note that `previous_len` rather than `len` is used + // for the filtering step because filtering is indeed for previous pass. + OffsetT previous_len; + + // We determine the bits of the k_th key inside the mask processed by the pass. The + // already known bits are stored in `kth_key_bits`. It's used to discriminate a + // element is a result (written to `out`), a candidate for next pass (written to + // `out_buf`), or not useful (discarded). The bits that are not yet processed do not + // matter for this purpose. + key_prefix_storage_t kth_key_bits; + + // Record how many elements have passed filtering. It's used to determine the position + // in the `out_buf` where an element should be written. + alignas(128) OffsetT filter_cnt; + + // For a row inside a batch, we may launch multiple thread blocks. This counter is + // used to determine if the current block is the last running block. If so, this block + // will execute compute_bin_offsets() and choose_bucket(). + alignas(128) unsigned int finished_block_cnt; + + // Record how many elements have been written to the front of `out`. Elements less (if + // SelectMin==true) than the k-th key are written from front to back. + alignas(128) OutOffsetT out_cnt; + + // Record how many elements have been written to the back of `out`. Elements equal to + // the k-th key are written from back to front. We need to keep count of them + // separately because the number of elements that <= the k-th key might exceed k. + alignas(128) OutOffsetT out_back_cnt; + // The 'alignas' is necessary to improve the performance of global memory accessing by isolating the request, + // especially for the segment version. +}; + +enum class candidate_class +{ + // The given candidate is definitely amongst the top-k items + selected, + // The given candidate may or may not be amongst the top-k items + candidate, + // The given candidate is definitely not amongst the top-k items + rejected +}; + +//! @brief AgentTopK implements a stateful abstraction of CUDA thread blocks for participating in +//! device-wide topK +//! +//! @tparam AgentTopKPolicyT +//! Parameterized agent_topk_policy tuning policy type +//! +//! @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 output iterator type for writing output values @iterator +//! +//! @tparam ExtractBinOpT +//! Operations to extract the bin from the input key values +//! +//! @tparam IdentifyCandidatesOpT +//! Operations to filter the input key values +//! +//! @tparam OffsetT +//! Type of variable num_items +//! +//! @tparam OutOffsetT +//! Type of variable k +//! +template +struct AgentTopK +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + // The key and value type + using key_in_t = it_value_t; + using value_in_t = it_value_t; + + static constexpr int threads_per_block = AgentTopKPolicyT::threads_per_block; + static constexpr int items_per_thread = AgentTopKPolicyT::items_per_thread; + static constexpr int bits_per_pass = AgentTopKPolicyT::bits_per_pass; + static constexpr int tile_items = threads_per_block * items_per_thread; + static constexpr int num_buckets = 1 << bits_per_pass; + + static constexpr bool keys_only = ::cuda::std::is_same_v; + static constexpr int bins_per_thread = ::cuda::ceil_div(num_buckets, threads_per_block); + + // Parameterized BlockLoad type for input data + using block_load_input_t = BlockLoad; + using block_load_trans_t = BlockLoad; + // Parameterized BlockScan type + using block_scan_t = BlockScan; + // Parameterized BlockStore type + using block_store_trans_t = BlockStore; + + // Shared memory + struct _TempStorage + { + union + { + // Smem needed for loading + typename block_load_input_t::TempStorage load_input; + typename block_load_trans_t::TempStorage load_trans; + // Smem needed for scan + typename block_scan_t::TempStorage scan; + // Smem needed for storing + typename block_store_trans_t::TempStorage store_trans; + }; + OffsetT histogram[num_buckets]; + }; + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + _TempStorage& temp_storage; // Reference to temp_storage + KeyInputIteratorT d_keys_in; // Input keys + KeyOutputIteratorT d_keys_out; // Output keys + ValueInputIteratorT d_values_in; // Input values + ValueOutputIteratorT d_values_out; // Output values + OffsetT num_items; // Total number of input items + OutOffsetT k; // Total number of output items + OffsetT buffer_length; // Size of the buffer for storing intermediate candidates + ExtractBinOpT extract_bin_op; // The operation for bin + IdentifyCandidatesOpT identify_candidates_op; // The operation for filtering + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + //! @param temp_storage + //! Reference to temp_storage + //! + //! @param d_keys_in + //! Input data, keys + //! + //! @param d_keys_out + //! Output data, keys + //! + //! @param d_values_in + //! Input data, values + //! + //! @param d_values_out + //! Output data, values + //! + //! @param num_items + //! Total number of input items + //! + //! @param k + //! The K value. Will find K elements from num_items elements + //! + //! @param buffer_length + //! The size of the buffer for storing intermediate candidates + //! + //! @param extract_bin_op + //! Extract bin operator + //! + //! @param identify_candidates_op + //! Filter operator + //! + _CCCL_DEVICE _CCCL_FORCEINLINE AgentTopK( + TempStorage& temp_storage, + const KeyInputIteratorT d_keys_in, + KeyOutputIteratorT d_keys_out, + const ValueInputIteratorT d_values_in, + ValueOutputIteratorT d_values_out, + OffsetT num_items, + OutOffsetT k, + OffsetT buffer_length, + ExtractBinOpT extract_bin_op, + IdentifyCandidatesOpT identify_candidates_op) + : temp_storage(temp_storage.Alias()) + , d_keys_in(d_keys_in) + , d_keys_out(d_keys_out) + , d_values_in(d_values_in) + , d_values_out(d_values_out) + , num_items(num_items) + , k(k) + , buffer_length(buffer_length) + , extract_bin_op(extract_bin_op) + , identify_candidates_op(identify_candidates_op) + {} + + //--------------------------------------------------------------------- + // Utility methods for device topK + //--------------------------------------------------------------------- + + // Process a range of input data in tiles, calling f(key, index) for each element + template + _CCCL_DEVICE _CCCL_FORCEINLINE void process_range(InputItT in, const OffsetT num_items, FuncT f) + { + key_in_t thread_data[items_per_thread]; + + const OffsetT items_per_pass = + static_cast(tile_items * gridDim.x); // NOLINT(bugprone-misplaced-widening-cast) + const OffsetT total_num_blocks = ::cuda::ceil_div(num_items, tile_items); + + const OffsetT num_remaining_elements = num_items % tile_items; + const OffsetT last_block_id = (total_num_blocks - 1) % gridDim.x; + + OffsetT tile_base = static_cast(blockIdx.x * tile_items); // NOLINT(bugprone-misplaced-widening-cast) + OffsetT offset = threadIdx.x * items_per_thread + tile_base; + + for (int i_block = static_cast(blockIdx.x); i_block < total_num_blocks - 1; + i_block += static_cast(gridDim.x)) + { + // Ensure that the temporary storage from previous iteration can be reused + __syncthreads(); + + block_load_input_t(temp_storage.load_input).Load(in + tile_base, thread_data); + for (int j = 0; j < items_per_thread; ++j) + { + f(thread_data[j], offset + j); + } + tile_base += items_per_pass; + offset += items_per_pass; + } + + // Last tile specialized code-path + if (blockIdx.x == last_block_id) + { + // Ensure that the temporary storage from the previous loop can be reused + __syncthreads(); + + if (num_remaining_elements == 0) + { + block_load_input_t(temp_storage.load_input).Load(in + tile_base, thread_data); + } + else + { + block_load_input_t(temp_storage.load_input).Load(in + tile_base, thread_data, num_remaining_elements); + } + + for (int j = 0; j < items_per_thread; ++j) + { + if ((offset + j) < num_items) + { + f(thread_data[j], offset + j); + } + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void init_histograms(OffsetT* histogram) + { + // Initialize histogram bin counts to zeros + int histo_offset = 0; + + // Loop unrolling is beneficial for performance here + _CCCL_PRAGMA_UNROLL_FULL() + for (; histo_offset + threads_per_block <= num_buckets; histo_offset += threads_per_block) + { + histogram[histo_offset + threadIdx.x] = 0; + } + // Finish up with guarded initialization if necessary + if ((num_buckets % threads_per_block != 0) && (histo_offset + threadIdx.x < num_buckets)) + { + histogram[histo_offset + threadIdx.x] = 0; + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void merge_histograms(OffsetT* global_histogram) + { + int histo_offset = 0; + + // Loop unrolling is beneficial for performance here + _CCCL_PRAGMA_UNROLL_FULL() + for (; histo_offset + threads_per_block <= num_buckets; histo_offset += threads_per_block) + { + if (temp_storage.histogram[histo_offset + threadIdx.x] != 0) + { + atomicAdd(global_histogram + (histo_offset + threadIdx.x), temp_storage.histogram[histo_offset + threadIdx.x]); + } + } + + // Finish up with guarded merging if necessary + if ((num_buckets % threads_per_block != 0) && (histo_offset + threadIdx.x < num_buckets)) + { + atomicAdd(global_histogram + (histo_offset + threadIdx.x), temp_storage.histogram[histo_offset + threadIdx.x]); + } + } + + // Fused filtering of the current pass and building histogram for the next pass + _CCCL_DEVICE _CCCL_FORCEINLINE void filter_and_histogram( + key_in_t* in_buf, + OffsetT* in_idx_buf, + key_in_t* out_buf, + OffsetT* out_idx_buf, + OffsetT previous_len, + Counter* counter, + OffsetT* histogram, + bool early_stop, + bool load_from_original_input) + { + // Initialize shared memory histogram + init_histograms(temp_storage.histogram); + + // Make sure the histogram was initialized + __syncthreads(); + + OffsetT* p_filter_cnt = &counter->filter_cnt; + OutOffsetT* p_out_cnt = &counter->out_cnt; + + // Lambda for early_stop = true (i.e., we have identified the exact "splitter" key): + // Select all items that fall into the bin of the k-th item (i.e., the 'candidates') and the ones that fall into + // bins preceding the k-th item bin (i.e., 'selected' items), write them to output. + // We can skip histogram computation because we don't need to further passes to refine the candidates. + auto f_early_stop = [load_from_original_input, in_idx_buf, p_out_cnt, this](key_in_t key, OffsetT i) { + const candidate_class pre_res = identify_candidates_op(key); + if (pre_res == candidate_class::candidate || pre_res == candidate_class::selected) + { + const OutOffsetT pos = atomicAdd(p_out_cnt, OutOffsetT{1}); + d_keys_out[pos] = key; + if constexpr (!keys_only) + { + const OffsetT index = load_from_original_input ? i : in_idx_buf[i]; + d_values_out[pos] = d_values_in[index]; + } + } + }; + + // Lambda for early_stop = false, out_buf != nullptr (i.e., we need to further refine the candidates in the next + // pass): Write out selected items to output, write candidates to out_buf, and build histogram for candidates. + auto f_with_out_buf = [load_from_original_input, in_idx_buf, out_buf, out_idx_buf, p_filter_cnt, p_out_cnt, this]( + key_in_t key, OffsetT i) { + const candidate_class pre_res = identify_candidates_op(key); + if (pre_res == candidate_class::candidate) + { + const OffsetT pos = atomicAdd(p_filter_cnt, OffsetT{1}); + out_buf[pos] = key; + if constexpr (!keys_only) + { + const OffsetT index = load_from_original_input ? i : in_idx_buf[i]; + out_idx_buf[pos] = index; + } + + const int bucket = extract_bin_op(key); + atomicAdd(temp_storage.histogram + bucket, OffsetT{1}); + } + else if (pre_res == candidate_class::selected) + { + const OutOffsetT pos = atomicAdd(p_out_cnt, OutOffsetT{1}); + d_keys_out[pos] = key; + if constexpr (!keys_only) + { + const OffsetT index = in_idx_buf ? in_idx_buf[i] : i; + d_values_out[pos] = d_values_in[index]; + } + } + }; + + // Lambda for early_stop = false, out_buf = nullptr (i.e., we need to further refine the candidates in the next + // pass, but we skip writing candidates to out_buf): + // Just build histogram for candidates. + // Note: We will only begin writing to d_keys_out starting from the pass in which the number of output-candidates + // is small enough to fit into the output buffer (otherwise, we would be writing the same items to d_keys_out + // multiple times). + auto f_no_out_buf = [this](key_in_t key, OffsetT i) { + const candidate_class pre_res = identify_candidates_op(key); + if (pre_res == candidate_class::candidate) + { + const int bucket = extract_bin_op(key); + atomicAdd(temp_storage.histogram + bucket, OffsetT{1}); + } + }; + + // Choose and invoke the appropriate lambda with the correct input source + // If the input size exceeds the allocated buffer size, we know for sure we haven't started writing candidates to + // the output buffer yet + if (load_from_original_input) + { + if (early_stop) + { + process_range(d_keys_in, previous_len, f_early_stop); + } + else if (out_buf) + { + process_range(d_keys_in, previous_len, f_with_out_buf); + } + else + { + process_range(d_keys_in, previous_len, f_no_out_buf); + } + } + else + { + if (early_stop) + { + process_range(in_buf, previous_len, f_early_stop); + } + else if (out_buf) + { + process_range(in_buf, previous_len, f_with_out_buf); + } + else + { + process_range(in_buf, previous_len, f_no_out_buf); + } + } + + // Early stop means that subsequent passes are not needed + if (early_stop) + { + return; + } + + // Ensure all threads have contributed to the histogram before accumulating in the global memory + __syncthreads(); + + // Merge the locally aggregated histogram into the global histogram + merge_histograms(histogram); + } + + // Replace histogram with its own prefix sum + _CCCL_DEVICE _CCCL_FORCEINLINE void compute_bin_offsets(volatile OffsetT* histogram) + { + OffsetT thread_data[bins_per_thread]{}; + + // Load global histogram (we can skip initializing oob-items to zero because they won't be stored back) + block_load_trans_t(temp_storage.load_trans).Load(histogram, thread_data, num_buckets); + __syncthreads(); + + block_scan_t(temp_storage.scan).InclusiveSum(thread_data, thread_data); + __syncthreads(); + + block_store_trans_t(temp_storage.store_trans).Store(temp_storage.histogram, thread_data, num_buckets); + } + + // Identify the bucket that the k-th value falls into + _CCCL_DEVICE _CCCL_FORCEINLINE void + choose_bucket(Counter* counter, const OutOffsetT k, const int pass) + { + // Initialize histogram bin counts to zeros + int histo_offset = 0; + + auto body = [&] { + const int bin_idx = static_cast(histo_offset + threadIdx.x); + const OffsetT prev = (bin_idx == 0) ? 0 : temp_storage.histogram[bin_idx - 1]; + const OffsetT cur = temp_storage.histogram[bin_idx]; + + // Identify the bin that the k-th item falls into. One and only one thread will satisfy this condition, so counter + // is written by only one thread + if (prev < k && cur >= k) + { + // The number of items that are yet to be identified + counter->k = k - prev; + + // The number of candidates in the next pass + counter->len = cur - prev; + const unsigned int bucket = static_cast(bin_idx); + // Update the "splitter" key by adding the radix digit of the k-th item bin of this pass + set_kth_key_bits(counter->kth_key_bits, pass, bucket); + } + }; + + _CCCL_PRAGMA_UNROLL_FULL() + for (; histo_offset + threads_per_block <= num_buckets; histo_offset += threads_per_block) + { + body(); + } + // Finish up with guarded initialization if necessary + if ((num_buckets % threads_per_block != 0) && (histo_offset + threadIdx.x < num_buckets)) + { + body(); + } + } + + // Performs the last-block coordination after histogram accumulation: ensures global visibility, + // detects the last finishing block, runs the prefix sum, identifies the k-th bucket, and resets + // the histogram for the next pass. The caller-supplied counter_update_fn runs on thread 0 of the + // last block to update pass-specific counter state. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void finalize_pass( + Counter* counter, + OffsetT* histogram, + OutOffsetT current_k, + int pass, + bool is_last_pass, + CounterUpdateFn counter_update_fn) + { + // Ensure all writes to the global memory-histogram are visible to all threads before + // proceeding to compute the prefix sum over the histogram. + __threadfence(); + + // Identify the last block in the grid to perform the prefix sum over the histogram + bool is_last_block = false; + if (threadIdx.x == 0) + { + unsigned int finished = atomicInc(&counter->finished_block_cnt, gridDim.x - 1); + is_last_block = (finished == (gridDim.x - 1)); + } + + // syncthreads ensures that the BlockLoad for loading the global histogram can reuse the temporary storage + if (__syncthreads_or(is_last_block)) + { + if (threadIdx.x == 0) + { + counter_update_fn(); + } + + // Compute prefix sum over the histogram's bin counts + compute_bin_offsets(histogram); + + // Make sure the prefix sum has been written to shared memory before choose_bucket() + __syncthreads(); + + // Identify the bucket that the k-th item falls into + choose_bucket(counter, current_k, pass); + + if (!is_last_pass) + { + init_histograms(histogram); + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void invoke_last_filter( + key_in_t* in_buf, OffsetT* in_idx_buf, Counter* counter, OutOffsetT k, int pass) + { + const bool load_from_original_input = (pass <= 1) || counter->previous_len > buffer_length; + const OffsetT current_len = load_from_original_input ? num_items : counter->previous_len; + in_idx_buf = load_from_original_input ? nullptr : in_idx_buf; // ? out_idx_buf : in_idx_buf; + + if (current_len == 0) + { + return; + } + + // changed in choose_bucket(); need to reload + OffsetT num_of_kth_needed = counter->k; + OutOffsetT* p_out_cnt = &counter->out_cnt; + OutOffsetT* p_out_back_cnt = &counter->out_back_cnt; + + auto f = [this, p_out_cnt, in_idx_buf, p_out_back_cnt, num_of_kth_needed, k, load_from_original_input]( + key_in_t key, OffsetT i) { + const candidate_class res = identify_candidates_op(key); + if (res == candidate_class::selected) + { + const OutOffsetT pos = atomicAdd(p_out_cnt, OffsetT{1}); + d_keys_out[pos] = key; + if constexpr (!keys_only) + { + // If writing has been skipped up to this point, `in_idx_buf` is nullptr + const OffsetT index = load_from_original_input ? i : in_idx_buf[i]; + d_values_out[pos] = d_values_in[index]; + } + } + else if (res == candidate_class::candidate) + { + const OutOffsetT back_pos = atomicAdd(p_out_back_cnt, OffsetT{1}); + + if (back_pos < num_of_kth_needed) + { + const OutOffsetT pos = k - 1 - back_pos; + d_keys_out[pos] = key; + if constexpr (!keys_only) + { + const OffsetT new_idx = load_from_original_input ? i : in_idx_buf[i]; + d_values_out[pos] = d_values_in[new_idx]; + } + } + } + }; + + if (load_from_original_input) + { + process_range(d_keys_in, current_len, f); + } + else + { + process_range(in_buf, current_len, f); + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void invoke_filter_and_histogram( + key_in_t* in_buf, + OffsetT* in_idx_buf, + key_in_t* out_buf, + OffsetT* out_idx_buf, + Counter* counter, + OffsetT* histogram, + int pass, + bool is_last_pass) + { + const OutOffsetT current_k = counter->k; + const OffsetT current_len = counter->len; + OffsetT previous_len = counter->previous_len; + + // If current_len is 0, it means all the candidates have been found in previous passes. + if (current_len == 0) + { + return; + } + + // Early stop means that the bin containing the k-th element has been identified, and all + // the elements in this bin are exactly the remaining k items we need to find. So we can + // stop the process after this filtering pass. + const bool early_stop = (current_len == static_cast(current_k)); + + // If previous_len > buffer_length, it means we haven't started writing candidates to out_buf yet, + // so have to make sure to load input directly from the original input. + // Also, unless we've had the chance to do at least one filtering pass, our input is definitely the original input + // (this is to guard against edge cases, e.g., buffer_length=num_items=1). + const bool load_from_original_input = (pass <= 1) || previous_len > buffer_length; + + if (load_from_original_input) + { + in_idx_buf = nullptr; + previous_len = num_items; + } + + // "current_len > buffer_length" means current pass will skip writing buffer + if (current_len > buffer_length) + { + out_buf = nullptr; + out_idx_buf = nullptr; + } + + // Fused filtering of candidates and histogram computation over the output-candidates + filter_and_histogram( + in_buf, in_idx_buf, out_buf, out_idx_buf, previous_len, counter, histogram, early_stop, load_from_original_input); + + finalize_pass(counter, histogram, current_k, pass, is_last_pass, [counter, current_len, early_stop] { + if (early_stop) + { + counter->previous_len = 0; + counter->len = 0; + } + else + { + counter->previous_len = current_len; + counter->filter_cnt = 0; + } + }); + } + + // Histogram-only pass: computes the histogram over the full input without filtering. + // Used for the first radix pass before any candidates have been identified. + _CCCL_DEVICE _CCCL_FORCEINLINE void invoke_histogram_only( + Counter* counter, OffsetT* histogram, int pass, bool is_last_pass) + { + // Initialize shared memory histogram + init_histograms(temp_storage.histogram); + __syncthreads(); + + // Compute per-thread block histograms over the full input + auto f = [this](key_in_t key, OffsetT /*index*/) { + const int bucket = extract_bin_op(key); + atomicAdd(temp_storage.histogram + bucket, OffsetT{1}); + }; + process_range(d_keys_in, num_items, f); + + // Ensure all threads have contributed to the histogram before accumulating in global memory + __syncthreads(); + + // Merge the locally aggregated histogram into the global histogram + merge_histograms(histogram); + + finalize_pass(counter, histogram, k, pass, is_last_pass, [counter, this] { + counter->previous_len = num_items; + counter->filter_cnt = 0; + }); + } +}; +} // namespace detail::topk +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_unique_by_key.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_unique_by_key.cuh new file mode 100644 index 00000000..a9ed003e --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/agent_unique_by_key.cuh @@ -0,0 +1,586 @@ +// SPDX-FileCopyrightText: Copyright (c), NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::AgentUniqueByKey implements a stateful abstraction of CUDA thread blocks for participating in device-wide + * unique-by-key. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Tuning policy types + ******************************************************************************/ + +namespace detail +{ +// TODO(bgruber): remove this when C++20 is the minimum, since then we can pass policy values as NTTP +template > +struct agent_unique_by_key_policy +{ + static constexpr int BLOCK_THREADS = ThreadsPerBlock; + static constexpr int ITEMS_PER_THREAD = ItemsPerThread; + static constexpr cub::BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm; + static constexpr cub::CacheLoadModifier LOAD_MODIFIER = LoadModifier; + static constexpr cub::BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm; + + struct detail + { + using delay_constructor_t = DelayConstructorT; + }; +}; +} // namespace detail + +//! Deprecated [Since 3.5] +template > +using AgentUniqueByKeyPolicy CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceSelect") = detail:: + agent_unique_by_key_policy; + +/****************************************************************************** + * Thread block abstractions + ******************************************************************************/ + +namespace detail::unique_by_key +{ +/** + * @brief AgentUniqueByKey implements a stateful abstraction of CUDA thread blocks for participating + * in device-wide unique-by-key + * + * @tparam AgentUniqueByKeyPolicyT + * Parameterized AgentUniqueByKeyPolicy tuning policy type + * + * @tparam KeyInputIteratorT + * Random-access input iterator type for keys + * + * @tparam ValueInputIteratorT + * Random-access input iterator type for values + * + * @tparam KeyOutputIteratorT + * Random-access output iterator type for keys + * + * @tparam ValueOutputIteratorT + * Random-access output iterator type for values + * + * @tparam EqualityOpT + * Equality operator type + * + * @tparam OffsetT + * Signed integer type for global offsets + */ +template +struct AgentUniqueByKey +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + // The input key and value type + using KeyT = cub::detail::it_value_t; + using ValueT = cub::detail::it_value_t; + + // Tile status descriptor interface type + using ScanTileStateT = ScanTileState; + + // Constants + static constexpr int BLOCK_THREADS = AgentUniqueByKeyPolicyT::BLOCK_THREADS; + static constexpr int ITEMS_PER_THREAD = AgentUniqueByKeyPolicyT::ITEMS_PER_THREAD; + static constexpr int ITEMS_PER_TILE = BLOCK_THREADS * ITEMS_PER_THREAD; + + // Cache-modified Input iterator wrapper type (for applying cache modifier) for keys + using WrappedKeyInputIteratorT = ::cuda::std::conditional_t< + ::cuda::std::is_pointer_v, + CacheModifiedInputIterator, // Wrap the native input pointer + // with + // CacheModifiedValuesInputIterator + KeyInputIteratorT>; // Directly use the supplied input iterator type + + // Cache-modified Input iterator wrapper type (for applying cache modifier) for values + using WrappedValueInputIteratorT = ::cuda::std::conditional_t< + ::cuda::std::is_pointer_v, + CacheModifiedInputIterator, // Wrap the native input + // pointer with + // CacheModifiedValuesInputIterator + ValueInputIteratorT>; // Directly use the supplied input iterator type + + // Parameterized BlockLoad type for input data + using BlockLoadKeys = BlockLoad; + + // Parameterized BlockLoad type for flags + using BlockLoadValues = BlockLoad; + + // Parameterized BlockDiscontinuity type for items + using BlockDiscontinuityKeys = cub::BlockDiscontinuity; + + // Parameterized BlockScan type + using BlockScanT = cub::BlockScan; + + // Parameterized BlockDiscontinuity type for items + using DelayConstructorT = typename AgentUniqueByKeyPolicyT::detail::delay_constructor_t; + using TilePrefixCallback = cub::TilePrefixCallbackOp, ScanTileStateT, DelayConstructorT>; + + // Key exchange type + using KeyExchangeT = KeyT[ITEMS_PER_TILE]; + + // Value exchange type + using ValueExchangeT = ValueT[ITEMS_PER_TILE]; + + // Shared memory type for this thread block + union _TempStorage + { + struct ScanStorage + { + typename BlockScanT::TempStorage scan; + typename TilePrefixCallback::TempStorage prefix; + typename BlockDiscontinuityKeys::TempStorage discontinuity; + } scan_storage; + + // Smem needed for loading keys + typename BlockLoadKeys::TempStorage load_keys; + + // Smem needed for loading values + typename BlockLoadValues::TempStorage load_values; + + // Smem needed for compacting items (allows non POD items in this union) + Uninitialized shared_keys; + Uninitialized shared_values; + }; + + // Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + _TempStorage& temp_storage; + WrappedKeyInputIteratorT d_keys_in; + WrappedValueInputIteratorT d_values_in; + KeyOutputIteratorT d_keys_out; + ValueOutputIteratorT d_values_out; + cub::InequalityWrapper inequality_op; + OffsetT num_items; + + //--------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------- + + // Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE AgentUniqueByKey( + TempStorage& temp_storage_, + WrappedKeyInputIteratorT d_keys_in_, + WrappedValueInputIteratorT d_values_in_, + KeyOutputIteratorT d_keys_out_, + ValueOutputIteratorT d_values_out_, + EqualityOpT equality_op_, + OffsetT num_items_) + : temp_storage(temp_storage_.Alias()) + , d_keys_in(d_keys_in_) + , d_values_in(d_values_in_) + , d_keys_out(d_keys_out_) + , d_values_out(d_values_out_) + , inequality_op(equality_op_) + , num_items(num_items_) + {} + + //--------------------------------------------------------------------- + // Utility functions + //--------------------------------------------------------------------- + + struct KeyTagT + {}; + struct ValueTagT + {}; + + _CCCL_DEVICE _CCCL_FORCEINLINE KeyExchangeT& GetShared(KeyTagT) + { + return temp_storage.shared_keys.Alias(); + } + _CCCL_DEVICE _CCCL_FORCEINLINE ValueExchangeT& GetShared(ValueTagT) + { + return temp_storage.shared_values.Alias(); + } + + //--------------------------------------------------------------------- + // Scatter utility methods + //--------------------------------------------------------------------- + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Scatter( + Tag tag, + OutputIt items_out, + T (&items)[ITEMS_PER_THREAD], + OffsetT (&selection_flags)[ITEMS_PER_THREAD], + OffsetT (&selection_indices)[ITEMS_PER_THREAD], + int /*num_tile_items*/, + int num_tile_selections, + OffsetT num_selections_prefix, + OffsetT /*num_selections*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + int local_scatter_offset = selection_indices[ITEM] - num_selections_prefix; + if (selection_flags[ITEM]) + { + GetShared(tag)[local_scatter_offset] = items[ITEM]; + } + } + + __syncthreads(); + + // Preventing loop unrolling helps avoid perf degradation when switching from signed to unsigned 32-bit offset + // types + _CCCL_PRAGMA_NOUNROLL() + for (int item = static_cast(threadIdx.x); item < num_tile_selections; item += BLOCK_THREADS) + { + items_out[num_selections_prefix + item] = GetShared(tag)[item]; // NOLINT(bugprone-misplaced-widening-cast) + } + + __syncthreads(); + } + + //--------------------------------------------------------------------- + // Cooperatively scan a device-wide sequence of tiles with other CTAs + //--------------------------------------------------------------------- + + /** + * @brief Process first tile of input (dynamic chained scan). + * + * @param num_tile_items + * Number of input items comprising this tile + * + * @param tile_offset + * Tile offset + * + * @param tile_state + * Global tile state descriptor + * + * @return The running count of selections (including this tile) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE OffsetT + ConsumeFirstTile(int num_tile_items, OffsetT tile_offset, ScanTileStateT& tile_state) + { + KeyT keys[ITEMS_PER_THREAD]; + OffsetT selection_flags[ITEMS_PER_THREAD]; + OffsetT selection_idx[ITEMS_PER_THREAD]; + + if constexpr (IS_LAST_TILE) + { + // Fill last elements with the first element + // because collectives are not suffix guarded + BlockLoadKeys(temp_storage.load_keys) + .Load(d_keys_in + tile_offset, keys, num_tile_items, *(d_keys_in + tile_offset)); + } + else + { + BlockLoadKeys(temp_storage.load_keys).Load(d_keys_in + tile_offset, keys); + } + + __syncthreads(); + + ValueT values[ITEMS_PER_THREAD]; + if constexpr (IS_LAST_TILE) + { + // Fill last elements with the first element + // because collectives are not suffix guarded + BlockLoadValues(temp_storage.load_values) + .Load(d_values_in + tile_offset, values, num_tile_items, *(d_values_in + tile_offset)); + } + else + { + BlockLoadValues(temp_storage.load_values).Load(d_values_in + tile_offset, values); + } + + __syncthreads(); + + BlockDiscontinuityKeys(temp_storage.scan_storage.discontinuity).FlagHeads(selection_flags, keys, inequality_op); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + // Set selection_flags for out-of-bounds items + if ((IS_LAST_TILE) && (OffsetT(threadIdx.x * ITEMS_PER_THREAD) + ITEM >= num_tile_items)) + { + selection_flags[ITEM] = 1; + } + } + + __syncthreads(); + + OffsetT num_tile_selections = 0; + OffsetT num_selections = 0; + OffsetT num_selections_prefix = 0; + + BlockScanT(temp_storage.scan_storage.scan).ExclusiveSum(selection_flags, selection_idx, num_tile_selections); + + if (threadIdx.x == 0) + { + // Update tile status if this is not the last tile + if constexpr (!IS_LAST_TILE) + { + tile_state.SetInclusive(0, num_tile_selections); + } + } + + // Do not count any out-of-bounds selections + if constexpr (IS_LAST_TILE) + { + int num_discount = ITEMS_PER_TILE - num_tile_items; + num_tile_selections -= num_discount; + } + num_selections = num_tile_selections; + + __syncthreads(); + + Scatter(KeyTagT(), + d_keys_out, + keys, + selection_flags, + selection_idx, + num_tile_items, + num_tile_selections, + num_selections_prefix, + num_selections); + + __syncthreads(); + + Scatter(ValueTagT(), + d_values_out, + values, + selection_flags, + selection_idx, + num_tile_items, + num_tile_selections, + num_selections_prefix, + num_selections); + + return num_selections; + } + + /** + * @brief Process subsequent tile of input (dynamic chained scan). + * + * @param num_tile_items + * Number of input items comprising this tile + * + * @param tile_idx + * Tile index + * + * @param tile_offset + * Tile offset + * + * @param tile_state + * Global tile state descriptor + * + * @return Returns the running count of selections (including this tile) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE OffsetT + ConsumeSubsequentTile(int num_tile_items, int tile_idx, OffsetT tile_offset, ScanTileStateT& tile_state) + { + KeyT keys[ITEMS_PER_THREAD]; + OffsetT selection_flags[ITEMS_PER_THREAD]; + OffsetT selection_idx[ITEMS_PER_THREAD]; + + if constexpr (IS_LAST_TILE) + { + // Fill last elements with the first element + // because collectives are not suffix guarded + BlockLoadKeys(temp_storage.load_keys) + .Load(d_keys_in + tile_offset, keys, num_tile_items, *(d_keys_in + tile_offset)); + } + else + { + BlockLoadKeys(temp_storage.load_keys).Load(d_keys_in + tile_offset, keys); + } + + __syncthreads(); + + ValueT values[ITEMS_PER_THREAD]; + if constexpr (IS_LAST_TILE) + { + // Fill last elements with the first element + // because collectives are not suffix guarded + BlockLoadValues(temp_storage.load_values) + .Load(d_values_in + tile_offset, values, num_tile_items, *(d_values_in + tile_offset)); + } + else + { + BlockLoadValues(temp_storage.load_values).Load(d_values_in + tile_offset, values); + } + + __syncthreads(); + + KeyT tile_predecessor = d_keys_in[tile_offset - 1]; + BlockDiscontinuityKeys(temp_storage.scan_storage.discontinuity) + .FlagHeads(selection_flags, keys, inequality_op, tile_predecessor); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ITEMS_PER_THREAD; ++ITEM) + { + // Set selection_flags for out-of-bounds items + if ((IS_LAST_TILE) && (OffsetT(threadIdx.x * ITEMS_PER_THREAD) + ITEM >= num_tile_items)) + { + selection_flags[ITEM] = 1; + } + } + + __syncthreads(); + + OffsetT num_tile_selections = 0; + OffsetT num_selections = 0; + OffsetT num_selections_prefix = 0; + + TilePrefixCallback prefix_cb(tile_state, temp_storage.scan_storage.prefix, ::cuda::std::plus<>{}, tile_idx); + BlockScanT(temp_storage.scan_storage.scan).ExclusiveSum(selection_flags, selection_idx, prefix_cb); + + num_selections = prefix_cb.GetInclusivePrefix(); + num_tile_selections = prefix_cb.GetBlockAggregate(); + num_selections_prefix = prefix_cb.GetExclusivePrefix(); + + if constexpr (IS_LAST_TILE) + { + int num_discount = ITEMS_PER_TILE - num_tile_items; + num_tile_selections -= num_discount; + num_selections -= num_discount; + } + + __syncthreads(); + + Scatter(KeyTagT(), + d_keys_out, + keys, + selection_flags, + selection_idx, + num_tile_items, + num_tile_selections, + num_selections_prefix, + num_selections); + + __syncthreads(); + + Scatter(ValueTagT(), + d_values_out, + values, + selection_flags, + selection_idx, + num_tile_items, + num_tile_selections, + num_selections_prefix, + num_selections); + + return num_selections; + } + + /** + * @brief Process a tile of input + * + * @param num_tile_items + * Number of input items comprising this tile + * + * @param tile_idx + * Tile index + * + * @param tile_offset + * Tile offset + * + * @param tile_state + * Global tile state descriptor + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE OffsetT + ConsumeTile(int num_tile_items, int tile_idx, OffsetT tile_offset, ScanTileStateT& tile_state) + { + OffsetT num_selections; + if (tile_idx == 0) + { + num_selections = ConsumeFirstTile(num_tile_items, tile_offset, tile_state); + } + else + { + num_selections = ConsumeSubsequentTile(num_tile_items, tile_idx, tile_offset, tile_state); + } + + return num_selections; + } + + /** + * @brief Scan tiles of items as part of a dynamic chained scan + * + * @param num_tiles + * Total number of input tiles + * + * @param tile_state + * Global tile state descriptor + * + * @param d_num_selected_out + * Output total number selection_flags + * + * @tparam NumSelectedIteratorT + * Output iterator type for recording number of items selection_flags + * + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ConsumeRange(int num_tiles, ScanTileStateT& tile_state, NumSelectedIteratorT d_num_selected_out) + { + // Blocks are launched in increasing order, so just assign one tile per block + int tile_idx = static_cast((blockIdx.x * gridDim.y) + blockIdx.y); // Current tile index + + // Global offset for the current tile + OffsetT tile_offset = static_cast(tile_idx) * static_cast(ITEMS_PER_TILE); + + if (tile_idx < num_tiles - 1) + { + ConsumeTile(ITEMS_PER_TILE, tile_idx, tile_offset, tile_state); + } + else + { + int num_remaining = static_cast(num_items - tile_offset); + OffsetT num_selections = ConsumeTile(num_remaining, tile_idx, tile_offset, tile_state); + if (threadIdx.x == 0) + { + *d_num_selected_out = num_selections; + } + } + } +}; +} // namespace detail::unique_by_key + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/agent/single_pass_scan_operators.cuh b/qwen3_6_scripts/cccl_preload/include/cub/agent/single_pass_scan_operators.cuh new file mode 100644 index 00000000..da336d5b --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/agent/single_pass_scan_operators.cuh @@ -0,0 +1,1391 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * Callback operator types for supplying BlockScan prefixes + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Prefix functor type for maintaining a running prefix while scanning a + * region independent of other thread blocks + ******************************************************************************/ + +/** + * Stateful callback operator type for supplying BlockScan prefixes. + * Maintains a running prefix that can be applied to consecutive + * BlockScan operations. + * + * @tparam T + * BlockScan value type + * + * @tparam ScanOpT + * Wrapped scan operator type + */ +template +struct BlockScanRunningPrefixOp +{ + /// Wrapped scan operator + ScanOpT op; + + /// Running block-wide prefix + T running_total; + + /// Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE BlockScanRunningPrefixOp(ScanOpT op) + : op(op) + {} + + /// Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE BlockScanRunningPrefixOp(T starting_prefix, ScanOpT op) + : op(op) + , running_total(starting_prefix) + {} + + /** + * Prefix callback operator. Returns the block-wide running_total in thread-0. + * + * @param block_aggregate + * The aggregate sum of the BlockScan inputs + */ + _CCCL_DEVICE _CCCL_FORCEINLINE T operator()(const T& block_aggregate) + { + T retval = running_total; + running_total = op(running_total, block_aggregate); + return retval; + } +}; + +/****************************************************************************** + * Generic tile status interface types for block-cooperative scans + ******************************************************************************/ + +/** + * Enumerations of tile status + */ +enum ScanTileStatus +{ + SCAN_TILE_OOB, // Out-of-bounds (e.g., padding) + SCAN_TILE_INVALID = 99, // Not yet processed + SCAN_TILE_PARTIAL, // Tile aggregate is available + SCAN_TILE_INCLUSIVE, // Inclusive tile prefix is available +}; + +/** + * Enum class used for specifying the memory order that shall be enforced while reading and writing the tile status. + */ +enum class MemoryOrder +{ + // Uses relaxed loads when reading a tile's status and relaxed stores when updating a tile's status + relaxed, + // Uses load acquire when reading a tile's status and store release when updating a tile's status + acquire_release +}; + +namespace detail +{ +template +_CCCL_DEVICE _CCCL_FORCEINLINE void delay() +{ + NV_IF_TARGET(NV_PROVIDES_SM_70, ({ + if (Delay > 0) + { + if (gridDim.x < GridThreshold) + { + __threadfence_block(); + } + else + { + __nanosleep(Delay); + } + } + })); +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void delay(int ns) +{ + NV_IF_TARGET(NV_PROVIDES_SM_70, ({ + if (ns > 0) + { + if (gridDim.x < GridThreshold) + { + __threadfence_block(); + } + else + { + __nanosleep(ns); + } + } + })); +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void always_delay() +{ + NV_IF_TARGET(NV_PROVIDES_SM_70, (__nanosleep(Delay);)); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void always_delay([[maybe_unused]] int ns) +{ + NV_IF_TARGET(NV_PROVIDES_SM_70, (__nanosleep(ns);)); +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void delay_or_prevent_hoisting() +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, (delay();), (__threadfence_block();)); +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void delay_or_prevent_hoisting([[maybe_unused]] int ns) +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, (delay(ns);), (__threadfence_block();)); +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void always_delay_or_prevent_hoisting() +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, (always_delay(Delay);), (__threadfence_block();)); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void always_delay_or_prevent_hoisting([[maybe_unused]] int ns) +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, (always_delay(ns);), (__threadfence_block();)); +} + +template +struct no_delay_constructor_t +{ + struct delay_t + { + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, (), (__threadfence_block();)); + } + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE no_delay_constructor_t(unsigned int /* seed */) + { + delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + return {}; + } +}; + +template +struct reduce_by_key_delay_constructor_t +{ + struct delay_t + { + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + NV_DISPATCH_TARGET( + NV_IS_EXACTLY_SM_80, + (delay();), + NV_PROVIDES_SM_70, + (delay<0, GridThreshold>();), + NV_IS_DEVICE, + (__threadfence_block();)); + } + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE reduce_by_key_delay_constructor_t(unsigned int /* seed */) + { + delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + return {}; + } +}; + +template +struct fixed_delay_constructor_t +{ + struct delay_t + { + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + delay_or_prevent_hoisting(); + } + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE fixed_delay_constructor_t(unsigned int /* seed */) + { + delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + return {}; + } +}; + +template +struct exponential_backoff_constructor_t +{ + struct delay_t + { + int delay; + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + always_delay_or_prevent_hoisting(delay); + delay <<= 1; + } + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE exponential_backoff_constructor_t(unsigned int /* seed */) + { + always_delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + return {InitialDelay}; + } +}; + +template +struct exponential_backoff_jitter_constructor_t +{ + struct delay_t + { + static constexpr unsigned int a = 16807; + static constexpr unsigned int c = 0; + static constexpr unsigned int m = 1u << 31; + + unsigned int max_delay; + unsigned int& seed; + + _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int next(unsigned int min, unsigned int max) + { + return (seed = (a * seed + c) % m) % (max + 1 - min) + min; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + always_delay_or_prevent_hoisting(next(0, max_delay)); + max_delay <<= 1; + } + }; + + unsigned int seed; + + _CCCL_DEVICE _CCCL_FORCEINLINE exponential_backoff_jitter_constructor_t(unsigned int seed) + : seed(seed) + { + always_delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + return {InitialDelay, seed}; + } +}; + +template +struct exponential_backoff_jitter_window_constructor_t +{ + struct delay_t + { + static constexpr unsigned int a = 16807; + static constexpr unsigned int c = 0; + static constexpr unsigned int m = 1u << 31; + + unsigned int max_delay; + unsigned int& seed; + + _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int next(unsigned int min, unsigned int max) + { + return (seed = (a * seed + c) % m) % (max + 1 - min) + min; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + unsigned int next_max_delay = max_delay << 1; + always_delay_or_prevent_hoisting(next(max_delay, next_max_delay)); + max_delay = next_max_delay; + } + }; + + unsigned int seed; + _CCCL_DEVICE _CCCL_FORCEINLINE exponential_backoff_jitter_window_constructor_t(unsigned int seed) + : seed(seed) + { + always_delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + return {InitialDelay, seed}; + } +}; + +template +struct exponential_backon_jitter_window_constructor_t +{ + struct delay_t + { + static constexpr unsigned int a = 16807; + static constexpr unsigned int c = 0; + static constexpr unsigned int m = 1u << 31; + + unsigned int max_delay; + unsigned int& seed; + + _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int next(unsigned int min, unsigned int max) + { + return (seed = (a * seed + c) % m) % (max + 1 - min) + min; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + int prev_delay = max_delay >> 1; + always_delay_or_prevent_hoisting(next(prev_delay, max_delay)); + max_delay = prev_delay; + } + }; + + unsigned int seed; + unsigned int max_delay = InitialDelay; + + _CCCL_DEVICE _CCCL_FORCEINLINE exponential_backon_jitter_window_constructor_t(unsigned int seed) + : seed(seed) + { + always_delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + max_delay >>= 1; + return {max_delay, seed}; + } +}; + +template +struct exponential_backon_jitter_constructor_t +{ + struct delay_t + { + static constexpr unsigned int a = 16807; + static constexpr unsigned int c = 0; + static constexpr unsigned int m = 1u << 31; + + unsigned int max_delay; + unsigned int& seed; + + _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int next(unsigned int min, unsigned int max) + { + return (seed = (a * seed + c) % m) % (max + 1 - min) + min; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + always_delay_or_prevent_hoisting(next(0, max_delay)); + max_delay >>= 1; + } + }; + + unsigned int seed; + unsigned int max_delay = InitialDelay; + + _CCCL_DEVICE _CCCL_FORCEINLINE exponential_backon_jitter_constructor_t(unsigned int seed) + : seed(seed) + { + always_delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + max_delay >>= 1; + return {max_delay, seed}; + } +}; + +template +struct exponential_backon_constructor_t +{ + struct delay_t + { + unsigned int delay; + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()() + { + always_delay_or_prevent_hoisting(delay); + delay >>= 1; + } + }; + + unsigned int max_delay = InitialDelay; + + _CCCL_DEVICE _CCCL_FORCEINLINE exponential_backon_constructor_t(unsigned int /* seed */) + { + always_delay(); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE delay_t operator()() + { + max_delay >>= 1; + return {max_delay}; + } +}; + +using default_no_delay_constructor_t = no_delay_constructor_t<450>; +using default_no_delay_t = default_no_delay_constructor_t::delay_t; + +template +using default_delay_constructor_t = + // TODO(bgruber): remove the check for is_primitive in CCCL 4.0 + ::cuda::std::conditional_t::value || ::cuda::is_trivially_copyable_v, + fixed_delay_constructor_t<350, 450>, + default_no_delay_constructor_t>; + +template +using default_delay_t = typename default_delay_constructor_t::delay_t; + +template +using default_reduce_by_key_delay_constructor_t = + // TODO(bgruber): remove the check for is_primitive in CCCL 4.0 + ::cuda::std::conditional_t<(is_primitive::value || ::cuda::is_trivially_copyable_v) + && (sizeof(ValueT) + sizeof(KeyT) < largest_atomic_message_size), + reduce_by_key_delay_constructor_t<350, 450>, + default_delay_constructor_t>>; + +/** + * @brief Alias template for a ScanTileState specialized for a given value type, `T`, and memory order `Order`. + * + * @tparam T The ScanTileState's value type + * @tparam Order The memory order to be implemented by the ScanTileState + */ +template +struct tile_state_with_memory_order +{ + ScanTileStateT& tile_state; + using T = typename ScanTileStateT::StatusValueT; + using StatusWord = typename ScanTileStateT::StatusWord; + + /** + * Update the specified tile's inclusive value and corresponding status + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void SetInclusive(int tile_idx, T tile_inclusive) + { + tile_state.template SetInclusive(tile_idx, tile_inclusive); + } + + /** + * Update the specified tile's partial value and corresponding status + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void SetPartial(int tile_idx, T tile_partial) + { + tile_state.template SetPartial(tile_idx, tile_partial); + } + + /** + * Wait for the corresponding tile to become non-invalid + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void WaitForValid(int tile_idx, StatusWord& status, T& value, DelayT delay = {}) + { + tile_state.template WaitForValid(tile_idx, status, value, delay); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE T LoadValid(int tile_idx) + { + return tile_state.template LoadValid(tile_idx); + } +}; + +_CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr size_t num_tiles_to_num_tile_states(size_t num_tiles) +{ + return warp_threads + num_tiles; +} + +_CCCL_HOST_DEVICE _CCCL_FORCEINLINE cudaError_t tile_state_allocation_size( + size_t& temp_storage_bytes, size_t bytes_per_description, size_t bytes_per_payload, size_t num_tiles) +{ + size_t num_tile_states = num_tiles_to_num_tile_states(num_tiles); + size_t allocation_sizes[]{ + // bytes needed for tile status descriptors + num_tile_states * bytes_per_description, + // bytes needed for partials + num_tile_states * bytes_per_payload, + // bytes needed for inclusives + num_tile_states * bytes_per_payload}; + // Set the necessary size of the blob + temp_storage_bytes = 0; + void* allocations[3] = {}; + return alias_temporaries(nullptr, temp_storage_bytes, allocations, allocation_sizes); +}; + +_CCCL_HOST_DEVICE _CCCL_FORCEINLINE cudaError_t tile_state_init( + size_t bytes_per_description, + size_t bytes_per_payload, + size_t num_tiles, + void* d_temp_storage, + size_t temp_storage_bytes, + void* (&allocations)[3]) +{ + size_t num_tile_states = num_tiles_to_num_tile_states(num_tiles); + size_t allocation_sizes[]{ + // bytes needed for tile status descriptors + num_tile_states * bytes_per_description, + // bytes needed for partials + num_tile_states * bytes_per_payload, + // bytes needed for inclusives + num_tile_states * bytes_per_payload}; + + // Set the necessary size of the blob + return alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes); +} +} // namespace detail + +/** + * Tile status interface. + */ +template in CCCL 4.0 + bool SingleWord = detail::is_primitive::value + || (::cuda::is_trivially_copyable_v + && sizeof(T) < detail::largest_atomic_message_size + // TODO(bgruber): a power of two size is not strictly necessary, but the implementation + // cannot handle it currently. For example, we could support status word + int3. + && ::cuda::is_power_of_two(sizeof(T)))> +struct ScanTileState; + +/** + * Tile status interface specialized for scan status and value types + * that can be combined into one machine word that can be + * read/written coherently in a single access. + */ +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document - causes Breathe/Sphinx parsing errors with nested templates +template +struct ScanTileState +{ + using StatusValueT = T; + + // Status word type + using StatusWord = ::cuda::std::_If< + sizeof(T) == 8, + unsigned long long, + ::cuda::std::_If>>; + + // Unit word type + using TxnWord = ::cuda::std::_If>; + static_assert(sizeof(TxnWord) <= detail::largest_atomic_message_size); + + // Device word type + struct TileDescriptor + { + StatusWord status; + T value; + }; + static_assert(sizeof(TileDescriptor) <= sizeof(TxnWord), "Tile descriptor must fit into the atomic transaction word"); + static_assert(sizeof(TileDescriptor) <= detail::largest_atomic_message_size); + + static constexpr int TILE_STATUS_PADDING = detail::warp_threads; + + // Device storage + TxnWord* d_tile_descriptors; + + static constexpr size_t description_bytes_per_tile = sizeof(TxnWord); + static constexpr size_t payload_bytes_per_tile = 0; + + /// Constructor + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE ScanTileState() + : d_tile_descriptors(nullptr) + {} + + /** + * @brief Initializer + * + * @param[in] d_temp_storage + * Device-accessible allocation of temporary storage. + */ + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE cudaError_t + Init(int /*num_tiles*/, void* d_temp_storage, size_t /*temp_storage_bytes*/) + { + d_tile_descriptors = reinterpret_cast(d_temp_storage); + return cudaSuccess; + } + + /** + * @brief Compute device memory needed for tile status + * + * @param[in] num_tiles + * Number of tiles + * + * @param[out] temp_storage_bytes + * Size in bytes of @p d_temp_storage allocation + */ + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static constexpr cudaError_t + AllocationSize(int num_tiles, size_t& temp_storage_bytes) + { + return detail::tile_state_allocation_size( + temp_storage_bytes, description_bytes_per_tile, payload_bytes_per_tile, num_tiles); + } + + /** + * Initialize (from device) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void InitializeStatus(int num_tiles) + { + int tile_idx = static_cast((blockIdx.x * blockDim.x) + threadIdx.x); + + TxnWord val = TxnWord(); + TileDescriptor* descriptor = reinterpret_cast(&val); + + if (tile_idx < num_tiles) + { + // Not-yet-set + descriptor->status = StatusWord(SCAN_TILE_INVALID); + d_tile_descriptors[TILE_STATUS_PADDING + tile_idx] = val; + } + + if ((blockIdx.x == 0) && (threadIdx.x < TILE_STATUS_PADDING)) + { + // Padding + descriptor->status = StatusWord(SCAN_TILE_OOB); + d_tile_descriptors[threadIdx.x] = val; + } + } + +private: + template + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::enable_if_t<(Order == MemoryOrder::relaxed), void> + StoreStatus(TxnWord* ptr, TxnWord alias) + { + detail::store_relaxed(ptr, alias); + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::enable_if_t<(Order == MemoryOrder::acquire_release), void> + StoreStatus(TxnWord* ptr, TxnWord alias) + { + detail::store_release(ptr, alias); + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::enable_if_t<(Order == MemoryOrder::relaxed), TxnWord> + LoadStatus(TxnWord* ptr) + { + return detail::load_relaxed(ptr); + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::enable_if_t<(Order == MemoryOrder::acquire_release), TxnWord> + LoadStatus(TxnWord* ptr) + { + // For pre-volta we hoist the memory barrier to outside the loop, i.e., after reading a valid state + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, (return detail::load_acquire(ptr);), (return detail::load_relaxed(ptr);)); + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::enable_if_t<(Order == MemoryOrder::relaxed), void> + ThreadfenceForLoadAcqPreVolta() + {} + + template + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::enable_if_t<(Order == MemoryOrder::acquire_release), void> + ThreadfenceForLoadAcqPreVolta() + { + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, (), (__threadfence();)); + } + +public: + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SetInclusive(int tile_idx, T tile_inclusive) + { + TileDescriptor tile_descriptor; + tile_descriptor.status = SCAN_TILE_INCLUSIVE; + tile_descriptor.value = tile_inclusive; + + TxnWord alias; + *reinterpret_cast(&alias) = tile_descriptor; + + StoreStatus(d_tile_descriptors + TILE_STATUS_PADDING + tile_idx, alias); + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SetPartial(int tile_idx, T tile_partial) + { + TileDescriptor tile_descriptor; + tile_descriptor.status = SCAN_TILE_PARTIAL; + tile_descriptor.value = tile_partial; + + TxnWord alias; + *reinterpret_cast(&alias) = tile_descriptor; + + StoreStatus(d_tile_descriptors + TILE_STATUS_PADDING + tile_idx, alias); + } + + /** + * Wait for the corresponding tile to become non-invalid + */ + template , MemoryOrder Order = MemoryOrder::relaxed> + _CCCL_DEVICE _CCCL_FORCEINLINE void + WaitForValid(int tile_idx, StatusWord& status, T& value, DelayT delay_or_prevent_hoisting = {}) + { + TileDescriptor tile_descriptor; + + { + TxnWord alias = LoadStatus(d_tile_descriptors + TILE_STATUS_PADDING + tile_idx); + tile_descriptor = reinterpret_cast(alias); + } + + while (__any_sync(0xffffffff, (tile_descriptor.status == SCAN_TILE_INVALID))) + { + delay_or_prevent_hoisting(); + TxnWord alias = LoadStatus(d_tile_descriptors + TILE_STATUS_PADDING + tile_idx); + tile_descriptor = reinterpret_cast(alias); + } + + // For pre-Volta and load acquire we emit relaxed loads in LoadStatus and hoist the threadfence here + ThreadfenceForLoadAcqPreVolta(); + + status = tile_descriptor.status; + value = tile_descriptor.value; + } + + /** + * Loads and returns the tile's value. The returned value is undefined if either (a) the tile's status is invalid or + * (b) there is no memory fence between reading a non-invalid status and the call to LoadValid. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE T LoadValid(int tile_idx) + { + TxnWord alias = d_tile_descriptors[TILE_STATUS_PADDING + tile_idx]; + TileDescriptor tile_descriptor = reinterpret_cast(alias); + return tile_descriptor.value; + } +}; +#endif // _CCCL_DOXYGEN_INVOKED + +/** + * Tile status interface specialized for scan status and value types that + * cannot be combined into one machine word. + */ +template +struct ScanTileState +{ + using StatusValueT = T; + + // Status word type + using StatusWord = unsigned int; + + static constexpr int TILE_STATUS_PADDING = detail::warp_threads; + + // Device storage + StatusWord* d_tile_status{}; + T* d_tile_partial{}; + T* d_tile_inclusive{}; + + static constexpr size_t description_bytes_per_tile = sizeof(StatusWord); + static constexpr size_t payload_bytes_per_tile = sizeof(Uninitialized); + + /// Constructor + _CCCL_FORCEINLINE ScanTileState() = default; + + /** + * @brief Initializer + * + * @param[in] num_tiles + * Number of tiles + * + * @param[in] d_temp_storage + * Device-accessible allocation of temporary storage. + * When nullptr, the required allocation size is written to \p temp_storage_bytes and no work is + * done. + * + * @param[in] temp_storage_bytes + * Size in bytes of @p d_temp_storage allocation + */ + /// Initializer + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE cudaError_t Init(int num_tiles, void* d_temp_storage, size_t temp_storage_bytes) + { + cudaError_t error = cudaSuccess; + do + { + void* allocations[3] = {}; + error = detail::tile_state_init( + description_bytes_per_tile, payload_bytes_per_tile, num_tiles, d_temp_storage, temp_storage_bytes, allocations); + if (cudaSuccess != error) + { + break; + } + // Alias the offsets + d_tile_status = reinterpret_cast(allocations[0]); + d_tile_partial = reinterpret_cast(allocations[1]); + d_tile_inclusive = reinterpret_cast(allocations[2]); + } while (false); + + return error; + } + + /** + * @brief Compute device memory needed for tile status + * + * @param[in] num_tiles + * Number of tiles + * + * @param[out] temp_storage_bytes + * Size in bytes of @p d_temp_storage allocation + */ + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static constexpr cudaError_t + AllocationSize(int num_tiles, size_t& temp_storage_bytes) + { + return detail::tile_state_allocation_size( + temp_storage_bytes, description_bytes_per_tile, payload_bytes_per_tile, num_tiles); + } + /** + * Initialize (from device) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void InitializeStatus(int num_tiles) + { + int tile_idx = static_cast((blockIdx.x * blockDim.x) + threadIdx.x); + if (tile_idx < num_tiles) + { + // Not-yet-set + d_tile_status[TILE_STATUS_PADDING + tile_idx] = StatusWord(SCAN_TILE_INVALID); + } + + if ((blockIdx.x == 0) && (threadIdx.x < TILE_STATUS_PADDING)) + { + // Padding + d_tile_status[threadIdx.x] = StatusWord(SCAN_TILE_OOB); + } + } + + /** + * Update the specified tile's inclusive value and corresponding status + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SetInclusive(int tile_idx, T tile_inclusive) + { + // Update tile inclusive value + ThreadStore(d_tile_inclusive + TILE_STATUS_PADDING + tile_idx, tile_inclusive); + detail::store_release(d_tile_status + TILE_STATUS_PADDING + tile_idx, StatusWord(SCAN_TILE_INCLUSIVE)); + } + + /** + * Update the specified tile's partial value and corresponding status + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SetPartial(int tile_idx, T tile_partial) + { + // Update tile partial value + ThreadStore(d_tile_partial + TILE_STATUS_PADDING + tile_idx, tile_partial); + detail::store_release(d_tile_status + TILE_STATUS_PADDING + tile_idx, StatusWord(SCAN_TILE_PARTIAL)); + } + + /** + * Wait for the corresponding tile to become non-invalid + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void WaitForValid(int tile_idx, StatusWord& status, T& value, DelayT delay = {}) + { + do + { + delay(); + status = detail::load_relaxed(d_tile_status + TILE_STATUS_PADDING + tile_idx); + __threadfence(); + } while (__any_sync(0xffffffff, (status == SCAN_TILE_INVALID))); + + if (status == StatusWord(SCAN_TILE_PARTIAL)) + { + value = ThreadLoad(d_tile_partial + TILE_STATUS_PADDING + tile_idx); + } + else if (status == StatusWord(SCAN_TILE_INCLUSIVE)) + { + value = ThreadLoad(d_tile_inclusive + TILE_STATUS_PADDING + tile_idx); + } + } + + /** + * Loads and returns the tile's value. The returned value is undefined if either (a) the tile's status is invalid or + * (b) there is no memory fence between reading a non-invalid status and the call to LoadValid. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE T LoadValid(int tile_idx) + { + return d_tile_inclusive[TILE_STATUS_PADDING + tile_idx]; + } +}; + +/****************************************************************************** + * ReduceByKey tile status interface types for block-cooperative scans + ******************************************************************************/ + +/** + * Tile status interface for reduction by key. + * + */ +template in CCCL 4.0 + bool SingleWord = (detail::is_primitive::value || ::cuda::is_trivially_copyable_v) + && (sizeof(ValueT) + sizeof(KeyT) < detail::largest_atomic_message_size)> +struct ReduceByKeyScanTileState; + +/** + * Tile status interface for reduction by key, specialized for scan status and value types that + * cannot be combined into one machine word. + */ +template +struct ReduceByKeyScanTileState : ScanTileState> +{ + using SuperClass = ScanTileState>; + + /// Constructor + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE ReduceByKeyScanTileState() + : SuperClass() + {} +}; + +/** + * Tile status interface for reduction by key, specialized for scan status and value types that + * can be combined into one machine word that can be read/written coherently in a single access. + */ +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document - causes Breathe/Sphinx parsing errors with nested templates +template +struct ReduceByKeyScanTileState +{ + using KeyValuePairT = KeyValuePair; + + // Constants + static constexpr int PAIR_SIZE = static_cast(sizeof(ValueT) + sizeof(KeyT)); + static constexpr int TXN_WORD_SIZE = 1 << Log2::VALUE; + static constexpr int STATUS_WORD_SIZE = TXN_WORD_SIZE - PAIR_SIZE; + + static constexpr int TILE_STATUS_PADDING = detail::warp_threads; + + // Status word type + using StatusWord = ::cuda::std::_If< + STATUS_WORD_SIZE == 8, + unsigned long long, + ::cuda::std:: + _If>>; + + // Status word type + using TxnWord = ::cuda::std:: + _If>; + + // Device word type (for when sizeof(ValueT) == sizeof(KeyT)) + struct TileDescriptorBigStatus + { + KeyT key; + ValueT value; + StatusWord status; + }; + + // Device word type (for when sizeof(ValueT) != sizeof(KeyT)) + struct TileDescriptorLittleStatus + { + ValueT value; + StatusWord status; + KeyT key; + }; + + // Device word type + using TileDescriptor = + ::cuda::std::_If; + + // Device storage + TxnWord* d_tile_descriptors; + + /// Constructor + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE ReduceByKeyScanTileState() + : d_tile_descriptors(nullptr) + {} + + /** + * @brief Initializer + * + * @param[in] num_tiles + * Number of tiles + * + * @param[in] d_temp_storage + * Device-accessible allocation of temporary storage. When nullptr, the required allocation size + * is written to \p temp_storage_bytes and no work is done. + * + * @param[in] temp_storage_bytes + * Size in bytes of @p d_temp_storage allocation + */ + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE cudaError_t + Init(int /*num_tiles*/, void* d_temp_storage, size_t /*temp_storage_bytes*/) + { + d_tile_descriptors = reinterpret_cast(d_temp_storage); + return cudaSuccess; + } + + /** + * @brief Compute device memory needed for tile status + * + * @param[in] num_tiles + * Number of tiles + * + * @param[out] temp_storage_bytes + * Size in bytes of @p d_temp_storage allocation + */ + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE static cudaError_t AllocationSize(int num_tiles, size_t& temp_storage_bytes) + { + // bytes needed for tile status descriptors + temp_storage_bytes = (num_tiles + TILE_STATUS_PADDING) * sizeof(TxnWord); + return cudaSuccess; + } + + /** + * Initialize (from device) + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void InitializeStatus(int num_tiles) + { + int tile_idx = static_cast((blockIdx.x * blockDim.x) + threadIdx.x); + TxnWord val = TxnWord(); + TileDescriptor* descriptor = reinterpret_cast(&val); + + if (tile_idx < num_tiles) + { + // Not-yet-set + descriptor->status = StatusWord(SCAN_TILE_INVALID); + d_tile_descriptors[TILE_STATUS_PADDING + tile_idx] = val; + } + + if ((blockIdx.x == 0) && (threadIdx.x < TILE_STATUS_PADDING)) + { + // Padding + descriptor->status = StatusWord(SCAN_TILE_OOB); + d_tile_descriptors[threadIdx.x] = val; + } + } + + /** + * Update the specified tile's inclusive value and corresponding status + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void SetInclusive(int tile_idx, KeyValuePairT tile_inclusive) + { + TileDescriptor tile_descriptor; + tile_descriptor.status = SCAN_TILE_INCLUSIVE; + tile_descriptor.value = tile_inclusive.value; + tile_descriptor.key = tile_inclusive.key; + + TxnWord alias; + *reinterpret_cast(&alias) = tile_descriptor; + + detail::store_relaxed(d_tile_descriptors + TILE_STATUS_PADDING + tile_idx, alias); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void SetPartial(int tile_idx, KeyValuePairT tile_partial) + { + TileDescriptor tile_descriptor; + tile_descriptor.status = SCAN_TILE_PARTIAL; + tile_descriptor.value = tile_partial.value; + tile_descriptor.key = tile_partial.key; + + TxnWord alias; + *reinterpret_cast(&alias) = tile_descriptor; + + detail::store_relaxed(d_tile_descriptors + TILE_STATUS_PADDING + tile_idx, alias); + } + + /** + * Wait for the corresponding tile to become non-invalid + */ + template ::delay_t> + _CCCL_DEVICE _CCCL_FORCEINLINE void + WaitForValid(int tile_idx, StatusWord& status, KeyValuePairT& value, DelayT delay_or_prevent_hoisting = {}) + { + // TxnWord alias = ThreadLoad(d_tile_descriptors + TILE_STATUS_PADDING + + // tile_idx); TileDescriptor tile_descriptor = reinterpret_cast(alias); + // + // while (tile_descriptor.status == SCAN_TILE_INVALID) + // { + // __threadfence_block(); // prevent hoisting loads from loop + // + // alias = ThreadLoad(d_tile_descriptors + TILE_STATUS_PADDING + tile_idx); + // tile_descriptor = reinterpret_cast(alias); + // } + // + // status = tile_descriptor.status; + // value.value = tile_descriptor.value; + // value.key = tile_descriptor.key; + + TileDescriptor tile_descriptor; + + do + { + delay_or_prevent_hoisting(); + TxnWord alias = detail::load_relaxed(d_tile_descriptors + TILE_STATUS_PADDING + tile_idx); + tile_descriptor = reinterpret_cast(alias); + + } while (__any_sync(0xffffffff, (tile_descriptor.status == SCAN_TILE_INVALID))); + + status = tile_descriptor.status; + value.value = tile_descriptor.value; + value.key = tile_descriptor.key; + } +}; +#endif // _CCCL_DOXYGEN_INVOKED + +/****************************************************************************** + * Prefix callback operator for coupling local block scan within a + * block-cooperative scan + ******************************************************************************/ + +/** + * Stateful block-scan prefix functor. Provides the running prefix for + * the current tile by using the callback warp to wait for + * aggregates/prefixes from predecessor tiles to become available. + * + * @tparam DelayConstructorT + * Implementation detail, do not specify directly, requirements on the + * content of this type are subject to breaking change. + */ +template , + bool StableReductionOrder = false> +struct TilePrefixCallbackOp +{ + // Parameterized warp reduce + using WarpReduceT = WarpReduce; + + // Temporary storage type + struct _TempStorage + { + typename WarpReduceT::TempStorage warp_reduce; + T exclusive_prefix; + T inclusive_prefix; + T block_aggregate; + }; + + // Alias wrapper allowing temporary storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + // Type of status word + using StatusWord = typename ScanTileStateT::StatusWord; + + // Fields + _TempStorage& temp_storage; ///< Reference to a warp-reduction instance + ScanTileStateT& tile_status; ///< Interface to tile status (non-anchor tiles stay at PARTIAL on the deterministic + ///< path) + ScanOpT scan_op; ///< Binary scan operator + int tile_idx; ///< The current tile index + T exclusive_prefix; ///< Exclusive prefix for the tile + T inclusive_prefix; ///< Inclusive prefix for the tile + + // Constructs prefix functor for a given tile index. + // Precondition: thread blocks processing all of the predecessor tiles were scheduled. + _CCCL_DEVICE _CCCL_FORCEINLINE + TilePrefixCallbackOp(ScanTileStateT& tile_status, TempStorage& temp_storage, ScanOpT scan_op, int tile_idx) + : temp_storage(temp_storage.Alias()) + , tile_status(tile_status) + , scan_op(scan_op) + , tile_idx(tile_idx) + {} + + // Computes the tile index and constructs prefix functor with it. + // Precondition: thread block per tile assignment. + _CCCL_DEVICE _CCCL_FORCEINLINE + TilePrefixCallbackOp(ScanTileStateT& tile_status, TempStorage& temp_storage, ScanOpT scan_op) + : TilePrefixCallbackOp(tile_status, temp_storage, scan_op, blockIdx.x) + {} + + /** + * @brief Block until all predecessors within the warp-wide window have non-invalid status + * + * @param predecessor_idx + * Preceding tile index to inspect + * + * @param[out] predecessor_status + * Preceding tile status + * + * @param[out] window_aggregate + * Relevant partial reduction from this window of preceding tiles + */ + template > + _CCCL_DEVICE _CCCL_FORCEINLINE void + ProcessWindow(int predecessor_idx, StatusWord& predecessor_status, T& window_aggregate, DelayT delay = {}) + { + T value; + tile_status.WaitForValid(predecessor_idx, predecessor_status, value, delay); + + // Perform a segmented reduction to get the prefix for the current window. + // Use the swizzled scan operator because we are now scanning *down* towards thread0. + + int tail_flag = (predecessor_status == StatusWord(SCAN_TILE_INCLUSIVE)); + window_aggregate = + WarpReduceT(temp_storage.warp_reduce).TailSegmentedReduce(value, tail_flag, SwizzleScanOp(scan_op)); + } + +private: + // Classic decoupled-lookback prefix computation. + _CCCL_DEVICE _CCCL_FORCEINLINE T lookback(T block_aggregate) + { + // Update our status with our tile-aggregate + if (threadIdx.x == 0) + { + detail::uninitialized_copy_single(&temp_storage.block_aggregate, block_aggregate); + + tile_status.SetPartial(tile_idx, block_aggregate); + } + + int predecessor_idx = tile_idx - threadIdx.x - 1; + StatusWord predecessor_status; + T window_aggregate; + + // Wait for the warp-wide window of predecessor tiles to become valid + DelayConstructorT construct_delay(tile_idx); + ProcessWindow(predecessor_idx, predecessor_status, window_aggregate, construct_delay()); + + // The exclusive tile prefix starts out as the current window aggregate + exclusive_prefix = window_aggregate; + + // Keep sliding the window back until we come across a tile whose inclusive prefix is known + while (__all_sync(0xffffffff, (predecessor_status != StatusWord(SCAN_TILE_INCLUSIVE)))) + { + predecessor_idx -= detail::warp_threads; + + // Update exclusive tile prefix with the window prefix + ProcessWindow(predecessor_idx, predecessor_status, window_aggregate, construct_delay()); + exclusive_prefix = scan_op(window_aggregate, exclusive_prefix); + } + + // Compute the inclusive tile prefix and update the status for this tile + if (threadIdx.x == 0) + { + inclusive_prefix = scan_op(exclusive_prefix, block_aggregate); + tile_status.SetInclusive(tile_idx, inclusive_prefix); + + detail::uninitialized_copy_single(&temp_storage.exclusive_prefix, exclusive_prefix); + + detail::uninitialized_copy_single(&temp_storage.inclusive_prefix, inclusive_prefix); + } + + // Return exclusive_prefix + return exclusive_prefix; + } + + // Run-to-run-deterministic K=1 32-batched lookback. Only anchor tiles publish INCLUSIVE. + _CCCL_DEVICE _CCCL_FORCEINLINE T lookback_stable_reduction_order(T block_aggregate) + { + if (threadIdx.x == 0) + { + detail::uninitialized_copy_single(&temp_storage.block_aggregate, block_aggregate); + tile_status.SetPartial(tile_idx, block_aggregate); + } + + const int predecessor_idx = tile_idx - threadIdx.x - 1; + StatusWord predecessor_status; + DelayConstructorT construct_delay(tile_idx); + + // lane that maps to the anchor tile (tile idx multiple of 32) + const int anchor_tile_lane = (tile_idx - 1) % detail::warp_threads; + + T value; + while (true) + { + tile_status.WaitForValid(predecessor_idx, predecessor_status, value, construct_delay()); + const int my_is_inclusive = (predecessor_status == StatusWord(SCAN_TILE_INCLUSIVE)); + const int anchor_is_inclusive = __shfl_sync(0xffffffff, my_is_inclusive, anchor_tile_lane); + if (anchor_is_inclusive) + { + break; + } + } + + const int tail_flag = (static_cast(threadIdx.x) == anchor_tile_lane); + exclusive_prefix = + WarpReduceT(temp_storage.warp_reduce).TailSegmentedReduce(value, tail_flag, SwizzleScanOp(scan_op)); + + if (threadIdx.x == 0) + { + inclusive_prefix = scan_op(exclusive_prefix, block_aggregate); + + // Only anchor tiles publish INCLUSIVE; non-anchor tiles stay at PARTIAL. + if (tile_idx % detail::warp_threads == 0) + { + tile_status.SetInclusive(tile_idx, inclusive_prefix); + } + + detail::uninitialized_copy_single(&temp_storage.exclusive_prefix, exclusive_prefix); + detail::uninitialized_copy_single(&temp_storage.inclusive_prefix, inclusive_prefix); + } + + return exclusive_prefix; + } + +public: + // BlockScan prefix callback functor. + _CCCL_DEVICE _CCCL_FORCEINLINE T operator()(T block_aggregate) + { + if constexpr (StableReductionOrder) + { + return lookback_stable_reduction_order(block_aggregate); + } + else + { + return lookback(block_aggregate); + } + } + + // Get the exclusive prefix stored in temporary storage + _CCCL_DEVICE _CCCL_FORCEINLINE T GetExclusivePrefix() + { + return temp_storage.exclusive_prefix; + } + + // Get the inclusive prefix stored in temporary storage + _CCCL_DEVICE _CCCL_FORCEINLINE T GetInclusivePrefix() + { + return temp_storage.inclusive_prefix; + } + + // Get the block aggregate stored in temporary storage + _CCCL_DEVICE _CCCL_FORCEINLINE T GetBlockAggregate() + { + return temp_storage.block_aggregate; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE int GetTileIdx() const + { + return tile_idx; + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_adjacent_difference.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_adjacent_difference.cuh new file mode 100644 index 00000000..e28588fa --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_adjacent_difference.cuh @@ -0,0 +1,971 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! The cub::BlockAdjacentDifference class provides collective methods for computing the differences of adjacent +//! elements partitioned across a CUDA thread block. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! BlockAdjacentDifference provides :ref:`collective ` methods for computing the +//! differences of adjacent elements partitioned across a CUDA thread block. +//! +//! Overview +//! ++++++++++++++++ +//! +//! BlockAdjacentDifference calculates the differences of adjacent elements in the elements partitioned across a CUDA +//! thread block. Because the binary operation could be noncommutative, there are two sets of methods. +//! Methods named SubtractLeft subtract left element ``i - 1`` of input sequence from current element ``i``. +//! Methods named SubtractRight subtract the right element ``i + 1`` from the current one ``i``: +//! +//! .. code-block:: c++ +//! +//! int values[4]; // [1, 2, 3, 4] +//! //... +//! int subtract_left_result[4]; <-- [ 1, 1, 1, 1 ] +//! int subtract_right_result[4]; <-- [ -1, -1, -1, 4 ] +//! +//! - For SubtractLeft, if the left element is out of bounds, the input value is assigned to ``output[0]`` +//! without modification. +//! - For SubtractRight, if the right element is out of bounds, the input value is assigned to the current output value +//! without modification. +//! - The block/example_block_reduce_dyn_smem.cu example under the examples/block folder illustrates usage of +//! dynamically shared memory with BlockReduce and how to re-purpose the same memory region. +//! This example can be easily adapted to the storage required by BlockAdjacentDifference. +//! +//! A Simple Example +//! ++++++++++++++++ +//! +//! The code snippet below illustrates how to use BlockAdjacentDifference to +//! compute the left difference between adjacent elements. +//! +//! .. code-block:: c++ +//! +//! #include +//! // or equivalently +//! +//! struct CustomDifference +//! { +//! template +//! __host__ DataType operator()(DataType &lhs, DataType &rhs) +//! { +//! return lhs - rhs; +//! } +//! }; +//! +//! __global__ void ExampleKernel(...) +//! { +//! // Specialize BlockAdjacentDifference for a 1D block of +//! // 128 threads of type int +//! using BlockAdjacentDifferenceT = +//! cub::BlockAdjacentDifference; +//! +//! // Allocate shared memory for BlockAdjacentDifference +//! __shared__ typename BlockAdjacentDifferenceT::TempStorage temp_storage; +//! +//! // Obtain a segment of consecutive items that are blocked across threads +//! int thread_data[4]; +//! ... +//! +//! // Collectively compute adjacent_difference +//! int result[4]; +//! +//! BlockAdjacentDifferenceT(temp_storage).SubtractLeft(thread_data, result, +//! CustomDifference()); +//! } +//! +//! Suppose the set of input `thread_data` across the block of threads is +//! ``{ [4,2,1,1], [1,1,1,1], [2,3,3,3], [3,4,1,4], ... }``. +//! The corresponding output ``result`` in those threads will be +//! ``{ [4,-2,-1,0], [0,0,0,0], [1,1,0,0], [0,1,-3,3], ... }``. +//! +//! @endrst +template +class BlockAdjacentDifference +{ +private: + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + /// Shared memory storage layout type (last element from each thread's input) + struct _TempStorage + { + T first_items[BLOCK_THREADS]; + T last_items[BLOCK_THREADS]; + }; + + /// Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + /// Specialization for when FlagOp has third index param + template ::value> + struct ApplyOp + { + // Apply flag operator + static _CCCL_DEVICE _CCCL_FORCEINLINE T FlagT(FlagOp flag_op, const T& a, const T& b, int idx) + { + return flag_op(b, a, idx); + } + }; + + /// Specialization for when FlagOp does not have a third index param + template + struct ApplyOp + { + // Apply flag operator + static _CCCL_DEVICE _CCCL_FORCEINLINE T FlagT(FlagOp flag_op, const T& a, const T& b, int /*idx*/) + { + return flag_op(b, a); + } + }; + + /// Templated unrolling of item comparison (inductive case) + struct Iterate + { + /** + * Head flags + * + * @param[out] flags Calling thread's discontinuity head_flags + * @param[in] input Calling thread's input items + * @param[out] preds Calling thread's predecessor items + * @param[in] flag_op Binary boolean flag predicate + */ + template + static _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeads( + int linear_tid, + FlagT (&flags)[ITEMS_PER_THREAD], + T (&input)[ITEMS_PER_THREAD], + T (&preds)[ITEMS_PER_THREAD], + FlagOp flag_op) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 1; i < ITEMS_PER_THREAD; ++i) + { + preds[i] = input[i - 1]; + flags[i] = ApplyOp::FlagT(flag_op, preds[i], input[i], (linear_tid * ITEMS_PER_THREAD) + i); + } + } + + /** + * Tail flags + * + * @param[out] flags Calling thread's discontinuity head_flags + * @param[in] input Calling thread's input items + * @param[in] flag_op Binary boolean flag predicate + */ + template + static _CCCL_DEVICE _CCCL_FORCEINLINE void + FlagTails(int linear_tid, FlagT (&flags)[ITEMS_PER_THREAD], T (&input)[ITEMS_PER_THREAD], FlagOp flag_op) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ITEMS_PER_THREAD - 1; ++i) + { + flags[i] = ApplyOp::FlagT(flag_op, input[i], input[i + 1], (linear_tid * ITEMS_PER_THREAD) + i + 1); + } + } + }; + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + unsigned int linear_tid; + +public: + /// @smemstorage{BlockAdjacentDifference} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using a private static allocation of shared memory as temporary storage + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + _CCCL_DEVICE _CCCL_FORCEINLINE BlockAdjacentDifference() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @brief Collective constructor using the specified memory allocation as temporary storage + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + //! + //! @param[in] temp_storage Reference to memory allocation having layout type TempStorage + _CCCL_DEVICE _CCCL_FORCEINLINE BlockAdjacentDifference(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Read left operations + //! @{ + + //! @rst + //! Subtracts the left element of each adjacent pair of elements partitioned across a CUDA thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates how to use BlockAdjacentDifference to compute the left difference between + //! adjacent elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockAdjacentDifference for a 1D block + //! // of 128 threads of type int + //! using BlockAdjacentDifferenceT = + //! cub::BlockAdjacentDifference; + //! + //! // Allocate shared memory for BlockAdjacentDifference + //! __shared__ typename BlockAdjacentDifferenceT::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively compute adjacent_difference + //! BlockAdjacentDifferenceT(temp_storage).SubtractLeft(thread_data, thread_data, + //! CustomDifference()); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [4,2,1,1], [1,1,1,1], [2,3,3,3], [3,4,1,4], ... }``. + //! The corresponding output ``result`` in those threads will be + //! ``{ [4,-2,-1,0], [0,0,0,0], [1,1,0,0], [0,1,-3,3], ... }``. + //! @endrst + //! + //! @param[out] output + //! Calling thread's adjacent difference result + //! + //! @param[in] input + //! Calling thread's input items (may be aliased to `output`) + //! + //! @param[in] difference_op + //! Binary difference operator + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + SubtractLeft(T (&input)[ITEMS_PER_THREAD], OutputType (&output)[ITEMS_PER_THREAD], DifferenceOpT difference_op) + { + // Share last item + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = ITEMS_PER_THREAD - 1; item > 0; item--) + { + output[item] = difference_op(input[item], input[item - 1]); + } + + if (linear_tid == 0) + { + output[0] = input[0]; + } + else + { + output[0] = difference_op(input[0], temp_storage.last_items[linear_tid - 1]); + } + } + + //! @rst + //! Subtracts the left element of each adjacent pair of elements partitioned across a CUDA thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates how to use BlockAdjacentDifference to compute the left difference between + //! adjacent elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockAdjacentDifference for a 1D block of + //! // 128 threads of type int + //! using BlockAdjacentDifferenceT = + //! cub::BlockAdjacentDifference; + //! + //! // Allocate shared memory for BlockAdjacentDifference + //! __shared__ typename BlockAdjacentDifferenceT::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // The last item in the previous tile: + //! int tile_predecessor_item = ...; + //! + //! // Collectively compute adjacent_difference + //! BlockAdjacentDifferenceT(temp_storage).SubtractLeft( + //! thread_data, + //! thread_data, + //! CustomDifference(), + //! tile_predecessor_item); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [4,2,1,1], [1,1,1,1], [2,3,3,3], [3,4,1,4], ... }``. + //! and that `tile_predecessor_item` is `3`. The corresponding output + //! ``result`` in those threads will be + //! ``{ [1,-2,-1,0], [0,0,0,0], [1,1,0,0], [0,1,-3,3], ... }``. + //! @endrst + //! + //! @param[out] output + //! Calling thread's adjacent difference result + //! + //! @param[in] input + //! Calling thread's input items (may be aliased to `output`) + //! + //! @param[in] difference_op + //! Binary difference operator + //! + //! @param[in] tile_predecessor_item + //! @rst + //! *thread*\ :sub:`0` only item which is going to be subtracted from the first tile item + //! (*input*\ :sub:`0` from *thread*\ :sub:`0`). + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SubtractLeft( + T (&input)[ITEMS_PER_THREAD], + OutputT (&output)[ITEMS_PER_THREAD], + DifferenceOpT difference_op, + T tile_predecessor_item) + { + // Share last item + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = ITEMS_PER_THREAD - 1; item > 0; item--) + { + output[item] = difference_op(input[item], input[item - 1]); + } + + // Set flag for first thread-item + if (linear_tid == 0) + { + output[0] = difference_op(input[0], tile_predecessor_item); + } + else + { + output[0] = difference_op(input[0], temp_storage.last_items[linear_tid - 1]); + } + } + + //! @rst + //! Subtracts the left element of each adjacent pair of elements partitioned across a CUDA thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates how to use BlockAdjacentDifference to compute the left difference between + //! adjacent elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockAdjacentDifference for a 1D block of + //! // 128 threads of type int + //! using BlockAdjacentDifferenceT = + //! cub::BlockAdjacentDifference; + //! + //! // Allocate shared memory for BlockAdjacentDifference + //! __shared__ typename BlockAdjacentDifferenceT::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! int valid_items = 9; + //! + //! // Collectively compute adjacent_difference + //! BlockAdjacentDifferenceT(temp_storage).SubtractLeftPartialTile( + //! thread_data, + //! thread_data, + //! CustomDifference(), + //! valid_items); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [4,2,1,1], [1,1,1,1], [2,3,3,3], [3,4,1,4], ... }``. + //! The corresponding output ``result`` in those threads will be + //! ``{ [4,-2,-1,0], [0,0,0,0], [1,3,3,3], [3,4,1,4], ... }``. + //! @endrst + //! + //! @param[out] output + //! Calling thread's adjacent difference result + //! + //! @param[in] input + //! Calling thread's input items (may be aliased to `output`) + //! + //! @param[in] difference_op + //! Binary difference operator + //! + //! @param[in] valid_items + //! Number of valid items in thread block + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SubtractLeftPartialTile( + T (&input)[ITEMS_PER_THREAD], OutputType (&output)[ITEMS_PER_THREAD], DifferenceOpT difference_op, int valid_items) + { + // Share last item + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + if ((linear_tid + 1) * ITEMS_PER_THREAD <= valid_items) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = ITEMS_PER_THREAD - 1; item > 0; item--) + { + output[item] = difference_op(input[item], input[item - 1]); + } + } + else + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = ITEMS_PER_THREAD - 1; item > 0; item--) + { + const int idx = linear_tid * ITEMS_PER_THREAD + item; + + if (idx < valid_items) + { + output[item] = difference_op(input[item], input[item - 1]); + } + else + { + output[item] = input[item]; + } + } + } + + if (linear_tid == 0 || valid_items <= linear_tid * ITEMS_PER_THREAD) + { + output[0] = input[0]; + } + else + { + output[0] = difference_op(input[0], temp_storage.last_items[linear_tid - 1]); + } + } + + //! @rst + //! Subtracts the left element of each adjacent pair of elements partitioned across a CUDA thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates how to use BlockAdjacentDifference to compute the left difference between + //! adjacent elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockAdjacentDifference for a 1D block of + //! // 128 threads of type int + //! using BlockAdjacentDifferenceT = + //! cub::BlockAdjacentDifference; + //! + //! // Allocate shared memory for BlockAdjacentDifference + //! __shared__ typename BlockAdjacentDifferenceT::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! int valid_items = 9; + //! int tile_predecessor_item = 4; + //! + //! // Collectively compute adjacent_difference + //! BlockAdjacentDifferenceT(temp_storage).SubtractLeftPartialTile( + //! thread_data, + //! thread_data, + //! CustomDifference(), + //! valid_items, + //! tile_predecessor_item); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [4,2,1,1], [1,1,1,1], [2,3,3,3], [3,4,1,4], ... }``. + //! The corresponding output ``result`` in those threads will be + //! ``{ [0,-2,-1,0], [0,0,0,0], [1,3,3,3], [3,4,1,4], ... }``. + //! @endrst + //! + //! @param[out] output + //! Calling thread's adjacent difference result + //! + //! @param[in] input + //! Calling thread's input items (may be aliased to `output`) + //! + //! @param[in] difference_op + //! Binary difference operator + //! + //! @param[in] valid_items + //! Number of valid items in thread block + //! + //! @param[in] tile_predecessor_item + //! @rst + //! *thread*\ :sub:`0` only item which is going to be subtracted from the first tile item + //! (*input*\ :sub:`0` from *thread*\ :sub:`0`). + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SubtractLeftPartialTile( + T (&input)[ITEMS_PER_THREAD], + OutputType (&output)[ITEMS_PER_THREAD], + DifferenceOpT difference_op, + int valid_items, + T tile_predecessor_item) + { + // Share last item + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + if ((linear_tid + 1) * ITEMS_PER_THREAD <= valid_items) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = ITEMS_PER_THREAD - 1; item > 0; item--) + { + output[item] = difference_op(input[item], input[item - 1]); + } + } + else + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = ITEMS_PER_THREAD - 1; item > 0; item--) + { + const int idx = linear_tid * ITEMS_PER_THREAD + item; + + if (idx < valid_items) + { + output[item] = difference_op(input[item], input[item - 1]); + } + else + { + output[item] = input[item]; + } + } + } + + if (valid_items <= linear_tid * ITEMS_PER_THREAD) + { + output[0] = input[0]; + } + else if (linear_tid == 0) + { + output[0] = difference_op(input[0], tile_predecessor_item); + } + else + { + output[0] = difference_op(input[0], temp_storage.last_items[linear_tid - 1]); + } + } + + //! @} + //! @name Read right operations + //! @{ + //! + //! @rst + //! + //! Subtracts the right element of each adjacent pair of elements partitioned across a CUDA thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates how to use BlockAdjacentDifference to compute the right difference between + //! adjacent elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockAdjacentDifference for a 1D block of + //! // 128 threads of type int + //! using BlockAdjacentDifferenceT = + //! cub::BlockAdjacentDifference; + //! + //! // Allocate shared memory for BlockAdjacentDifference + //! __shared__ typename BlockAdjacentDifferenceT::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively compute adjacent_difference + //! BlockAdjacentDifferenceT(temp_storage).SubtractRight( + //! thread_data, + //! thread_data, + //! CustomDifference()); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ ...3], [4,2,1,1], [1,1,1,1], [2,3,3,3], [3,4,1,4] }``. + //! The corresponding output ``result`` in those threads will be + //! ``{ ...-1, [2,1,0,0], [0,0,0,-1], [-1,0,0,0], [-1,3,-3,4] }``. + //! @endrst + //! + //! @param[out] output + //! Calling thread's adjacent difference result + //! + //! @param[in] input + //! Calling thread's input items (may be aliased to `output`) + //! + //! @param[in] difference_op + //! Binary difference operator + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + SubtractRight(T (&input)[ITEMS_PER_THREAD], OutputT (&output)[ITEMS_PER_THREAD], DifferenceOpT difference_op) + { + // Share first item + temp_storage.first_items[linear_tid] = input[0]; + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ITEMS_PER_THREAD - 1; item++) + { + output[item] = difference_op(input[item], input[item + 1]); + } + + if (linear_tid == BLOCK_THREADS - 1) + { + output[ITEMS_PER_THREAD - 1] = input[ITEMS_PER_THREAD - 1]; + } + else + { + output[ITEMS_PER_THREAD - 1] = + difference_op(input[ITEMS_PER_THREAD - 1], temp_storage.first_items[linear_tid + 1]); + } + } + + //! @rst + //! Subtracts the right element of each adjacent pair of elements partitioned across a CUDA thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates how to use BlockAdjacentDifference to compute the right difference between + //! adjacent elements. + //! + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockAdjacentDifference for a 1D block of + //! // 128 threads of type int + //! using BlockAdjacentDifferenceT = + //! cub::BlockAdjacentDifference; + //! + //! // Allocate shared memory for BlockAdjacentDifference + //! __shared__ typename BlockAdjacentDifferenceT::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // The first item in the next tile: + //! int tile_successor_item = ...; + //! + //! // Collectively compute adjacent_difference + //! BlockAdjacentDifferenceT(temp_storage).SubtractRight( + //! thread_data, + //! thread_data, + //! CustomDifference(), + //! tile_successor_item); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ ...3], [4,2,1,1], [1,1,1,1], [2,3,3,3], [3,4,1,4] }``, + //! and that ``tile_successor_item`` is ``3``. The corresponding output ``result`` + //! in those threads will be + //! ``{ ...-1, [2,1,0,0], [0,0,0,-1], [-1,0,0,0], [-1,3,-3,1] }``. + //! @endrst + //! + //! @param[out] output + //! Calling thread's adjacent difference result + //! + //! @param[in] input + //! Calling thread's input items (may be aliased to `output`) + //! + //! @param[in] difference_op + //! Binary difference operator + //! + //! @param[in] tile_successor_item + //! @rst + //! *thread*\ :sub:`BLOCK_THREADS` only item which is going to be subtracted from the last tile item + //! (*input*\ :sub:`ITEMS_PER_THREAD` from *thread*\ :sub:`BLOCK_THREADS`). + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SubtractRight( + T (&input)[ITEMS_PER_THREAD], + OutputT (&output)[ITEMS_PER_THREAD], + DifferenceOpT difference_op, + T tile_successor_item) + { + // Share first item + temp_storage.first_items[linear_tid] = input[0]; + + __syncthreads(); + + // Set flag for last thread-item + T successor_item = (linear_tid == BLOCK_THREADS - 1) + ? tile_successor_item // Last thread + : temp_storage.first_items[linear_tid + 1]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ITEMS_PER_THREAD - 1; item++) + { + output[item] = difference_op(input[item], input[item + 1]); + } + + output[ITEMS_PER_THREAD - 1] = difference_op(input[ITEMS_PER_THREAD - 1], successor_item); + } + + //! @rst + //! Subtracts the right element of each adjacent pair in range of elements partitioned across a CUDA thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates how to use BlockAdjacentDifference to compute the right difference between + //! adjacent elements. + //! + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockAdjacentDifference for a 1D block of + //! // 128 threads of type int + //! using BlockAdjacentDifferenceT = + //! cub::BlockAdjacentDifference; + //! + //! // Allocate shared memory for BlockAdjacentDifference + //! __shared__ typename BlockAdjacentDifferenceT::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively compute adjacent_difference + //! BlockAdjacentDifferenceT(temp_storage).SubtractRightPartialTile( + //! thread_data, + //! thread_data, + //! CustomDifference(), + //! valid_items); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ ...3], [4,2,1,1], [1,1,1,1], [2,3,3,3], [3,4,1,4] }``. + //! and that ``valid_items`` is ``507``. The corresponding output ``result`` in + //! those threads will be + //! ``{ ...-1, [2,1,0,0], [0,0,0,-1], [-1,0,3,3], [3,4,1,4] }``. + //! @endrst + //! + //! @param[out] output + //! Calling thread's adjacent difference result + //! + //! @param[in] input + //! Calling thread's input items (may be aliased to `output`) + //! + //! @param[in] difference_op + //! Binary difference operator + //! + //! @param[in] valid_items + //! Number of valid items in thread block + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SubtractRightPartialTile( + T (&input)[ITEMS_PER_THREAD], OutputT (&output)[ITEMS_PER_THREAD], DifferenceOpT difference_op, int valid_items) + { + // Share first item + temp_storage.first_items[linear_tid] = input[0]; + + __syncthreads(); + + if ((linear_tid + 1) * ITEMS_PER_THREAD < valid_items) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ITEMS_PER_THREAD - 1; item++) + { + output[item] = difference_op(input[item], input[item + 1]); + } + + output[ITEMS_PER_THREAD - 1] = + difference_op(input[ITEMS_PER_THREAD - 1], temp_storage.first_items[linear_tid + 1]); + } + else + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ITEMS_PER_THREAD; item++) + { + const int idx = linear_tid * ITEMS_PER_THREAD + item; + + // Right element of input[valid_items - 1] is out of bounds. + // According to the API it's copied into output array + // without modification. + if (idx < valid_items - 1) + { + output[item] = difference_op(input[item], input[item + 1]); + } + else + { + output[item] = input[item]; + } + } + } + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_discontinuity.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_discontinuity.cuh new file mode 100644 index 00000000..fbe57d97 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_discontinuity.cuh @@ -0,0 +1,1234 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * The cub::BlockDiscontinuity class provides [collective](../index.html#sec0) methods for + * flagging discontinuities within an ordered set of items partitioned across a CUDA thread block. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! The BlockDiscontinuity class provides :ref:`collective ` methods for +//! flagging discontinuities within an ordered set of items partitioned across a CUDA thread +//! block. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - A set of "head flags" (or "tail flags") is often used to indicate corresponding items +//! that differ from their predecessors (or successors). For example, head flags are convenient +//! for demarcating disjoint data segments as part of a segmented scan or reduction. +//! - @blocked +//! +//! Performance Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - @granularity +//! - Incurs zero bank conflicts for most types +//! +//! A Simple Example +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @blockcollective{BlockDiscontinuity} +//! +//! The code snippet below illustrates the head flagging of 512 integer items that +//! are partitioned in a :ref:`blocked arrangement ` across 128 threads +//! where each thread owns 4 consecutive items. +//! +//! .. code-block:: c++ +//! +//! #include // or equivalently +//! +//! __global__ void ExampleKernel(...) +//! { +//! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int +//! using BlockDiscontinuity = cub::BlockDiscontinuity; +//! +//! // Allocate shared memory for BlockDiscontinuity +//! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; +//! +//! // Obtain a segment of consecutive items that are blocked across threads +//! int thread_data[4]; +//! ... +//! +//! // Collectively compute head flags for discontinuities in the segment +//! int head_flags[4]; +//! BlockDiscontinuity(temp_storage).FlagHeads(head_flags, thread_data, cub::Inequality()); +//! } +//! +//! Suppose the set of input ``thread_data`` across the block of threads is +//! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], [3,4,4,4], ... }``. +//! The corresponding output ``head_flags`` in those threads will be +//! ``{ [1,0,1,0], [0,0,0,0], [1,1,0,0], [0,1,0,0], ... }``. +//! +//! Re-using dynamically allocating shared memory +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``examples/block/example_block_reduce_dyn_smem.cu`` example illustrates usage of +//! dynamically shared memory with BlockReduce and how to re-purpose the same memory region. +//! This example can be easily adapted to the storage required by BlockDiscontinuity. +//! @endrst +//! +//! @tparam T +//! The data type to be flagged. +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockDiscontinuity +{ +private: + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + /// Shared memory storage layout type (last element from each thread's input) + struct _TempStorage + { + T first_items[BLOCK_THREADS]; + T last_items[BLOCK_THREADS]; + }; + + /// Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + /// Specialization for when FlagOp has third index param + template ::value> + struct ApplyOp + { + // Apply flag operator + static _CCCL_DEVICE _CCCL_FORCEINLINE bool FlagT(FlagOp flag_op, const T& a, const T& b, int idx) + { + return flag_op(a, b, idx); + } + }; + + /// Specialization for when FlagOp does not have a third index param + template + struct ApplyOp + { + // Apply flag operator + static _CCCL_DEVICE _CCCL_FORCEINLINE bool FlagT(FlagOp flag_op, const T& a, const T& b, int /*idx*/) + { + return flag_op(a, b); + } + }; + + /// Templated unrolling of item comparison (inductive case) + struct Iterate + { + /** + * @brief Head flags + * + * @param[out] flags + * Calling thread's discontinuity head_flags + * + * @param[in] input + * Calling thread's input items + * + * @param[out] preds + * Calling thread's predecessor items + * + * @param[in] flag_op + * Binary boolean flag predicate + */ + template + static _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeads( + int linear_tid, + FlagT (&flags)[ITEMS_PER_THREAD], + T (&input)[ITEMS_PER_THREAD], + T (&preds)[ITEMS_PER_THREAD], + FlagOp flag_op) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 1; i < ITEMS_PER_THREAD; ++i) + { + preds[i] = input[i - 1]; + flags[i] = ApplyOp::FlagT(flag_op, preds[i], input[i], (linear_tid * ITEMS_PER_THREAD) + i); + } + } + + /** + * @brief Tail flags + * + * @param[out] flags + * Calling thread's discontinuity head_flags + * + * @param[in] input + * Calling thread's input items + * + * @param[in] flag_op + * Binary boolean flag predicate + */ + template + static _CCCL_DEVICE _CCCL_FORCEINLINE void + FlagTails(int linear_tid, FlagT (&flags)[ITEMS_PER_THREAD], T (&input)[ITEMS_PER_THREAD], FlagOp flag_op) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ITEMS_PER_THREAD - 1; ++i) + { + flags[i] = ApplyOp::FlagT(flag_op, input[i], input[i + 1], (linear_tid * ITEMS_PER_THREAD) + i + 1); + } + } + }; + + /****************************************************************************** + * Thread fields + ******************************************************************************/ + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + unsigned int linear_tid; + +public: + /// @smemstorage{BlockDiscontinuity} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + /** + * @brief Collective constructor using a private static allocation of shared memory as temporary + * storage. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockDiscontinuity() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @brief Collective constructor using the specified memory allocation as temporary storage. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @param[in] temp_storage + * Reference to memory allocation having layout type TempStorage + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockDiscontinuity(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Head flag operations + //! @{ + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + + /** + * @param[out] head_flags + * Calling thread's discontinuity head_flags + * + * @param[in] input + * Calling thread's input items + * + * @param[out] preds + * Calling thread's predecessor items + * + * @param[in] flag_op + * Binary boolean flag predicate + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeads( + FlagT (&head_flags)[ITEMS_PER_THREAD], T (&input)[ITEMS_PER_THREAD], T (&preds)[ITEMS_PER_THREAD], FlagOp flag_op) + { + // Share last item + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + if (linear_tid == 0) + { + // Set flag for first thread-item (preds[0] is undefined) + head_flags[0] = 1; + } + else + { + preds[0] = temp_storage.last_items[linear_tid - 1]; + head_flags[0] = ApplyOp::FlagT(flag_op, preds[0], input[0], linear_tid * ITEMS_PER_THREAD); + } + + // Set head_flags for remaining items + Iterate::FlagHeads(linear_tid, head_flags, input, preds, flag_op); + } + + /** + * @param[out] head_flags + * Calling thread's discontinuity head_flags + * + * @param[in] input + * Calling thread's input items + * + * @param[out] preds + * Calling thread's predecessor items + * + * @param[in] flag_op + * Binary boolean flag predicate + * + * @param[in] tile_predecessor_item + * [thread0 only] Item with which to compare the first tile item + * (input0 from thread0). + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeads( + FlagT (&head_flags)[ITEMS_PER_THREAD], + T (&input)[ITEMS_PER_THREAD], + T (&preds)[ITEMS_PER_THREAD], + FlagOp flag_op, + T tile_predecessor_item) + { + // Share last item + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + // Set flag for first thread-item + preds[0] = (linear_tid == 0) ? tile_predecessor_item : // First thread + temp_storage.last_items[linear_tid - 1]; + + head_flags[0] = ApplyOp::FlagT(flag_op, preds[0], input[0], linear_tid * ITEMS_PER_THREAD); + + // Set head_flags for remaining items + Iterate::FlagHeads(linear_tid, head_flags, input, preds, flag_op); + } + +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Sets head flags indicating discontinuities between items partitioned across the thread + //! block, for which the first item has no reference and is always flagged. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The flag ``head_flags[i]`` is set for item ``input[i]`` when ``flag_op(previous-item, input[i])`` returns + //! ``true`` (where ``previous-item`` is either the preceding item in the same thread or the last item in + //! the previous thread). + //! - For *thread*\ :sub:`0`, item ``input[0]`` is always flagged. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the head-flagging of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int + //! using BlockDiscontinuity = cub::BlockDiscontinuity; + //! + //! // Allocate shared memory for BlockDiscontinuity + //! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively compute head flags for discontinuities in the segment + //! int head_flags[4]; + //! BlockDiscontinuity(temp_storage).FlagHeads(head_flags, thread_data, cub::Inequality()); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], [3,4,4,4], ... }``. + //! The corresponding output ``head_flags`` in those threads will be + //! ``{ [1,0,1,0], [0,0,0,0], [1,1,0,0], [0,1,0,0], ... }``. + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread + //! + //! @tparam FlagT + //! **[inferred]** The flag type (must be an integer type) + //! + //! @tparam FlagOp + //! **[inferred]** Binary predicate functor type having member + //! `T operator()(const T &a, const T &b)` or member + //! `T operator()(const T &a, const T &b, unsigned int b_index)`, and returning `true` + //! if a discontinuity exists between `a` and `b`, otherwise `false`. + //! `b_index` is the rank of b in the aggregate tile of data. + //! + //! @param[out] head_flags + //! Calling thread's discontinuity head_flags + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[in] flag_op + //! Binary boolean flag predicate + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + FlagHeads(FlagT (&head_flags)[ITEMS_PER_THREAD], T (&input)[ITEMS_PER_THREAD], FlagOp flag_op) + { + T preds[ITEMS_PER_THREAD]; + FlagHeads(head_flags, input, preds, flag_op); + } + + //! @rst + //! Sets head flags indicating discontinuities between items partitioned across the thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The flag ``head_flags[i]`` is set for item ``input[i]`` when ``flag_op(previous-item, input[i])`` + //! returns ``true`` (where ``previous-item`` is either the preceding item in the same thread or the last item + //! in the previous thread). + //! - For *thread*\ :sub:`0`, item ``input[0]`` is compared against ``tile_predecessor_item``. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the head-flagging of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int + //! using BlockDiscontinuity = cub::BlockDiscontinuity; + //! + //! // Allocate shared memory for BlockDiscontinuity + //! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Have thread0 obtain the predecessor item for the entire tile + //! int tile_predecessor_item; + //! if (threadIdx.x == 0) tile_predecessor_item == ... + //! + //! // Collectively compute head flags for discontinuities in the segment + //! int head_flags[4]; + //! BlockDiscontinuity(temp_storage).FlagHeads(head_flags, thread_data, + //! cub::Inequality(), tile_predecessor_item); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], [3,4,4,4], ... }``, + //! and that ``tile_predecessor_item`` is ``0``. The corresponding output ``head_flags`` in those + //! threads will be ``{ [0,0,1,0], [0,0,0,0], [1,1,0,0], [0,1,0,0], ... }``. + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam FlagT + //! **[inferred]** The flag type (must be an integer type) + //! + //! @tparam FlagOp + //! **[inferred]** Binary predicate functor type having member + //! `T operator()(const T &a, const T &b)` or member + //! `T operator()(const T &a, const T &b, unsigned int b_index)`, + //! and returning `true` if a discontinuity exists between `a` and `b`, + //! otherwise `false`. `b_index` is the rank of b in the aggregate tile of data. + //! + //! @param[out] head_flags + //! Calling thread's discontinuity `head_flags` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[in] flag_op + //! Binary boolean flag predicate + //! + //! @param[in] tile_predecessor_item + //! @rst + //! *thread*\ :sub:`0` only item with which to compare the first tile item (``input[0]`` from *thread*\ :sub:`0`). + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeads( + FlagT (&head_flags)[ITEMS_PER_THREAD], T (&input)[ITEMS_PER_THREAD], FlagOp flag_op, T tile_predecessor_item) + { + T preds[ITEMS_PER_THREAD]; + FlagHeads(head_flags, input, preds, flag_op, tile_predecessor_item); + } + + //! @} + //! @name Tail flag operations + //! @{ + + //! @rst + //! Sets tail flags indicating discontinuities between items partitioned across the thread + //! block, for which the last item has no reference and is always flagged. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The flag ``tail_flags[i]`` is set for item ``input[i]`` when + //! ``flag_op(input[i], next-item)`` + //! returns ``true`` (where `next-item` is either the next item + //! in the same thread or the first item in the next thread). + //! - For *thread*\ :sub:`BLOCK_THREADS - 1`, item ``input[ITEMS_PER_THREAD - 1]`` is always flagged. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the tail-flagging of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int + //! using BlockDiscontinuity = cub::BlockDiscontinuity; + //! + //! // Allocate shared memory for BlockDiscontinuity + //! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively compute tail flags for discontinuities in the segment + //! int tail_flags[4]; + //! BlockDiscontinuity(temp_storage).FlagTails(tail_flags, thread_data, cub::Inequality()); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], ..., [124,125,125,125] }``. + //! The corresponding output ``tail_flags`` in those threads will be + //! ``{ [0,1,0,0], [0,0,0,1], [1,0,0,...], ..., [1,0,0,1] }``. + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam FlagT + //! **[inferred]** The flag type (must be an integer type) + //! + //! @tparam FlagOp + //! **[inferred]** Binary predicate functor type having member + //! `T operator()(const T &a, const T &b)` or member + //! `T operator()(const T &a, const T &b, unsigned int b_index)`, and returning `true` + //! if a discontinuity exists between `a` and `b`, otherwise `false`. `b_index` is the + //! rank of `b` in the aggregate tile of data. + //! + //! @param[out] tail_flags + //! Calling thread's discontinuity tail_flags + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[in] flag_op + //! Binary boolean flag predicate + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + FlagTails(FlagT (&tail_flags)[ITEMS_PER_THREAD], T (&input)[ITEMS_PER_THREAD], FlagOp flag_op) + { + // Share first item + temp_storage.first_items[linear_tid] = input[0]; + + __syncthreads(); + + // Set flag for last thread-item + tail_flags[ITEMS_PER_THREAD - 1] = + (linear_tid == BLOCK_THREADS - 1) ? 1 : // Last thread + ApplyOp::FlagT( + flag_op, + input[ITEMS_PER_THREAD - 1], + temp_storage.first_items[linear_tid + 1], + (linear_tid * ITEMS_PER_THREAD) + ITEMS_PER_THREAD); + + // Set tail_flags for remaining items + Iterate::FlagTails(linear_tid, tail_flags, input, flag_op); + } + + //! @rst + //! Sets tail flags indicating discontinuities between items partitioned across the thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The flag ``tail_flags[i]`` is set for item ``input[i]`` when ``flag_op(input[i], next-item)`` + //! returns ``true`` (where ``next-item`` is either the next item in the same thread or the first item in + //! the next thread). + //! - For *thread*\ :sub:`BLOCK_THREADS - 1`, item ``input[ITEMS_PER_THREAD - 1]`` is compared against + //! ``tile_successor_item``. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the tail-flagging of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int + //! using BlockDiscontinuity = cub::BlockDiscontinuity; + //! + //! // Allocate shared memory for BlockDiscontinuity + //! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Have thread127 obtain the successor item for the entire tile + //! int tile_successor_item; + //! if (threadIdx.x == 127) tile_successor_item == ... + //! + //! // Collectively compute tail flags for discontinuities in the segment + //! int tail_flags[4]; + //! BlockDiscontinuity(temp_storage).FlagTails(tail_flags, thread_data, + //! cub::Inequality(), tile_successor_item); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], ..., [124,125,125,125] }`` + //! and that ``tile_successor_item`` is ``125``. The corresponding output ``tail_flags`` in those + //! threads will be ``{ [0,1,0,0], [0,0,0,1], [1,0,0,...], ..., [1,0,0,0] }``. + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam FlagT + //! **[inferred]** The flag type (must be an integer type) + //! + //! @tparam FlagOp + //! **[inferred]** Binary predicate functor type having member + //! `T operator()(const T &a, const T &b)` or member + //! `T operator()(const T &a, const T &b, unsigned int b_index)`, and returning `true` + //! if a discontinuity exists between `a` and `b`, otherwise `false`. `b_index` is the + //! rank of `b` in the aggregate tile of data. + //! + //! @param[out] tail_flags + //! Calling thread's discontinuity tail_flags + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[in] flag_op + //! Binary boolean flag predicate + //! + //! @param[in] tile_successor_item + //! @rst + //! *thread*\ :sub:`BLOCK_THREADS - 1` only item with which to + //! compare the last tile item (``input[ITEMS_PER_THREAD - 1]`` from + //! *thread*\ :sub:`BLOCK_THREADS - 1`). + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + FlagTails(FlagT (&tail_flags)[ITEMS_PER_THREAD], T (&input)[ITEMS_PER_THREAD], FlagOp flag_op, T tile_successor_item) + { + // Share first item + temp_storage.first_items[linear_tid] = input[0]; + + __syncthreads(); + + // Set flag for last thread-item + T successor_item = (linear_tid == BLOCK_THREADS - 1) ? tile_successor_item : // Last thread + temp_storage.first_items[linear_tid + 1]; + + tail_flags[ITEMS_PER_THREAD - 1] = ApplyOp::FlagT( + flag_op, input[ITEMS_PER_THREAD - 1], successor_item, (linear_tid * ITEMS_PER_THREAD) + ITEMS_PER_THREAD); + + // Set tail_flags for remaining items + Iterate::FlagTails(linear_tid, tail_flags, input, flag_op); + } + + //! @} + //! @name Head & tail flag operations + //! @{ + + //! @rst + //! Sets both head and tail flags indicating discontinuities between items partitioned across the thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The flag ``head_flags[i]`` is set for item ``input[i]`` when ``flag_op(previous-item, input[i])`` returns + //! ``true`` (where ``previous-item`` is either the preceding item in the same thread or the last item in + //! the previous thread). + //! - For *thread*\ :sub:`0`, item ``input[0]`` is always flagged. + //! - The flag ``tail_flags[i]`` is set for item ``input[i]`` when ``flag_op(input[i], next-item)`` + //! returns ``true`` (where next-item is either the next item in the same thread or the first item in + //! the next thread). + //! - For *thread*\ :sub:`BLOCK_THREADS - 1`, item ``input[ITEMS_PER_THREAD - 1]`` is always flagged. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the head- and tail-flagging of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int + //! using BlockDiscontinuity = cub::BlockDiscontinuity; + //! + //! // Allocate shared memory for BlockDiscontinuity + //! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively compute head and flags for discontinuities in the segment + //! int head_flags[4]; + //! int tail_flags[4]; + //! BlockDiscontinuity(temp_storage).FlagHeadsAndTails(head_flags, tail_flags, thread_data, + //! cub::Inequality()); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], ..., [124,125,125,125] }`` + //! and that the tile_successor_item is ``125``. The corresponding output ``head_flags`` + //! in those threads will be ``{ [1,0,1,0], [0,0,0,0], [1,1,0,0], [0,1,0,0], ... }``. + //! and the corresponding output ``tail_flags`` in those threads will be + //! ``{ [0,1,0,0], [0,0,0,1], [1,0,0,...], ..., [1,0,0,1] }``. + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam FlagT + //! **[inferred]** The flag type (must be an integer type) + //! + //! @tparam FlagOp + //! **[inferred]** Binary predicate functor type having member + //! `T operator()(const T &a, const T &b)` or member + //! `T operator()(const T &a, const T &b, unsigned int b_index)`, and returning `true` + //! if a discontinuity exists between `a` and `b`, otherwise `false`. `b_index` is the + //! rank of `b` in the aggregate tile of data. + //! + //! @param[out] head_flags + //! Calling thread's discontinuity head_flags + //! + //! @param[out] tail_flags + //! Calling thread's discontinuity tail_flags + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[in] flag_op + //! Binary boolean flag predicate + template + _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeadsAndTails( + FlagT (&head_flags)[ITEMS_PER_THREAD], + FlagT (&tail_flags)[ITEMS_PER_THREAD], + T (&input)[ITEMS_PER_THREAD], + FlagOp flag_op) + { + // Share first and last items + temp_storage.first_items[linear_tid] = input[0]; + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + T preds[ITEMS_PER_THREAD]; + + // Set flag for first thread-item + if (linear_tid == 0) + { + head_flags[0] = 1; + } + else + { + preds[0] = temp_storage.last_items[linear_tid - 1]; + head_flags[0] = ApplyOp::FlagT(flag_op, preds[0], input[0], linear_tid * ITEMS_PER_THREAD); + } + + // Set flag for last thread-item + tail_flags[ITEMS_PER_THREAD - 1] = + (linear_tid == BLOCK_THREADS - 1) ? 1 : // Last thread + ApplyOp::FlagT( + flag_op, + input[ITEMS_PER_THREAD - 1], + temp_storage.first_items[linear_tid + 1], + (linear_tid * ITEMS_PER_THREAD) + ITEMS_PER_THREAD); + + // Set head_flags for remaining items + Iterate::FlagHeads(linear_tid, head_flags, input, preds, flag_op); + + // Set tail_flags for remaining items + Iterate::FlagTails(linear_tid, tail_flags, input, flag_op); + } + + //! @rst + //! Sets both head and tail flags indicating discontinuities between items partitioned across the thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The flag ``head_flags[i]`` is set for item ``input[i]`` when + //! ``flag_op(previous-item, input[i])`` returns ``true`` (where ``previous-item`` is either the preceding item + //! in the same thread or the last item in the previous thread). + //! - For *thread*\ :sub:`0`, item ``input[0]`` is always flagged. + //! - The flag ``tail_flags[i]`` is set for item ``input[i]`` when ``flag_op(input[i], next-item)`` returns ``true`` + //! (where ``next-item`` is either the next item in the same thread or the first item in the next thread). + //! - For *thread*\ :sub:`BLOCK_THREADS - 1`, item ``input[ITEMS_PER_THREAD - 1]`` is compared + //! against ``tile_predecessor_item``. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the head- and tail-flagging of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int + //! using BlockDiscontinuity = cub::BlockDiscontinuity; + //! + //! // Allocate shared memory for BlockDiscontinuity + //! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Have thread127 obtain the successor item for the entire tile + //! int tile_successor_item; + //! if (threadIdx.x == 127) tile_successor_item == ... + //! + //! // Collectively compute head and flags for discontinuities in the segment + //! int head_flags[4]; + //! int tail_flags[4]; + //! BlockDiscontinuity(temp_storage).FlagHeadsAndTails(head_flags, tail_flags, + //! tile_successor_item, thread_data, + //! cub::Inequality()); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], ..., [124,125,125,125] }`` + //! and that the tile_successor_item is ``125``. The corresponding output ``head_flags`` + //! in those threads will be ``{ [1,0,1,0], [0,0,0,0], [1,1,0,0], [0,1,0,0], ... }``. + //! and the corresponding output ``tail_flags`` in those threads will be + //! ``{ [0,1,0,0], [0,0,0,1], [1,0,0,...], ..., [1,0,0,0] }``. + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam FlagT + //! **[inferred]** The flag type (must be an integer type) + //! + //! @tparam FlagOp + //! **[inferred]** Binary predicate functor type having member + //! `T operator()(const T &a, const T &b)` or member + //! `T operator()(const T &a, const T &b, unsigned int b_index)`, and returning `true` + //! if a discontinuity exists between `a` and `b`, otherwise `false`. `b_index` is the + //! rank of b in the aggregate tile of data. + //! + //! @param[out] head_flags + //! Calling thread's discontinuity head_flags + //! + //! @param[out] tail_flags + //! Calling thread's discontinuity tail_flags + //! + //! @param[in] tile_successor_item + //! @rst + //! *thread*\ :sub:`BLOCK_THREADS - 1` only item with which to compare + //! the last tile item (``input[ITEMS_PER_THREAD - 1]`` from + //! *thread*\ :sub:`BLOCK_THREADS - 1`). + //! @endrst + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[in] flag_op + //! Binary boolean flag predicate + template + _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeadsAndTails( + FlagT (&head_flags)[ITEMS_PER_THREAD], + FlagT (&tail_flags)[ITEMS_PER_THREAD], + T tile_successor_item, + T (&input)[ITEMS_PER_THREAD], + FlagOp flag_op) + { + // Share first and last items + temp_storage.first_items[linear_tid] = input[0]; + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + T preds[ITEMS_PER_THREAD]; + + // Set flag for first thread-item + if (linear_tid == 0) + { + head_flags[0] = 1; + } + else + { + preds[0] = temp_storage.last_items[linear_tid - 1]; + head_flags[0] = ApplyOp::FlagT(flag_op, preds[0], input[0], linear_tid * ITEMS_PER_THREAD); + } + + // Set flag for last thread-item + T successor_item = (linear_tid == BLOCK_THREADS - 1) ? tile_successor_item : // Last thread + temp_storage.first_items[linear_tid + 1]; + + tail_flags[ITEMS_PER_THREAD - 1] = ApplyOp::FlagT( + flag_op, input[ITEMS_PER_THREAD - 1], successor_item, (linear_tid * ITEMS_PER_THREAD) + ITEMS_PER_THREAD); + + // Set head_flags for remaining items + Iterate::FlagHeads(linear_tid, head_flags, input, preds, flag_op); + + // Set tail_flags for remaining items + Iterate::FlagTails(linear_tid, tail_flags, input, flag_op); + } + + //! @rst + //! Sets both head and tail flags indicating discontinuities between items partitioned across the thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The flag ``head_flags[i]`` is set for item ``input[i]`` when ``flag_op(previous-item, input[i])`` + //! returns ``true`` (where ``previous-item`` is either the preceding item in the same thread or the last item + //! in the previous thread). + //! - For *thread*\ :sub:`0`, item ``input[0]`` is compared against ``tile_predecessor_item``. + //! - The flag ``tail_flags[i]`` is set for item ``input[i]`` when + //! ``flag_op(input[i], next-item)`` returns ``true`` (where ``next-item`` is either the next item + //! in the same thread or the first item in the next thread). + //! - For *thread*\ :sub:`BLOCK_THREADS - 1`, item + //! ``input[ITEMS_PER_THREAD - 1]`` is always flagged. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the head- and tail-flagging of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int + //! using BlockDiscontinuity = cub::BlockDiscontinuity; + //! + //! // Allocate shared memory for BlockDiscontinuity + //! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Have thread0 obtain the predecessor item for the entire tile + //! int tile_predecessor_item; + //! if (threadIdx.x == 0) tile_predecessor_item == ... + //! + //! // Have thread127 obtain the successor item for the entire tile + //! int tile_successor_item; + //! if (threadIdx.x == 127) tile_successor_item == ... + //! + //! // Collectively compute head and flags for discontinuities in the segment + //! int head_flags[4]; + //! int tail_flags[4]; + //! BlockDiscontinuity(temp_storage).FlagHeadsAndTails(head_flags, tile_predecessor_item, + //! tail_flags, tile_successor_item, + //! thread_data, cub::Inequality()); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], ..., [124,125,125,125] }``, + //! that the ``tile_predecessor_item`` is ``0``, and that the ``tile_successor_item`` is ``125``. + //! The corresponding output ``head_flags`` in those threads will be + //! ``{ [0,0,1,0], [0,0,0,0], [1,1,0,0], [0,1,0,0], ... }``, and the corresponding output ``tail_flags`` + //! in those threads will be ``{ [0,1,0,0], [0,0,0,1], [1,0,0,...], ..., [1,0,0,1] }``. + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam FlagT + //! **[inferred]** The flag type (must be an integer type) + //! + //! @tparam FlagOp + //! **[inferred]** Binary predicate functor type having member + //! `T operator()(const T &a, const T &b)` or member + //! `T operator()(const T &a, const T &b, unsigned int b_index)`, and returning `true` + //! if a discontinuity exists between `a` and `b`, otherwise `false`. `b_index` is the rank + //! of b in the aggregate tile of data. + //! + //! @param[out] head_flags + //! Calling thread's discontinuity head_flags + //! + //! @param[in] tile_predecessor_item + //! @rst + //! *thread*\ :sub:`0` only item with which to compare the first tile item (``input[0]`` from *thread*\ :sub:`0`). + //! @endrst + //! + //! @param[out] tail_flags + //! Calling thread's discontinuity tail_flags + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[in] flag_op + //! Binary boolean flag predicate + template + _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeadsAndTails( + FlagT (&head_flags)[ITEMS_PER_THREAD], + T tile_predecessor_item, + FlagT (&tail_flags)[ITEMS_PER_THREAD], + T (&input)[ITEMS_PER_THREAD], + FlagOp flag_op) + { + // Share first and last items + temp_storage.first_items[linear_tid] = input[0]; + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + T preds[ITEMS_PER_THREAD]; + + // Set flag for first thread-item + preds[0] = (linear_tid == 0) ? tile_predecessor_item : // First thread + temp_storage.last_items[linear_tid - 1]; + + head_flags[0] = ApplyOp::FlagT(flag_op, preds[0], input[0], linear_tid * ITEMS_PER_THREAD); + + // Set flag for last thread-item + tail_flags[ITEMS_PER_THREAD - 1] = + (linear_tid == BLOCK_THREADS - 1) ? 1 : // Last thread + ApplyOp::FlagT( + flag_op, + input[ITEMS_PER_THREAD - 1], + temp_storage.first_items[linear_tid + 1], + (linear_tid * ITEMS_PER_THREAD) + ITEMS_PER_THREAD); + + // Set head_flags for remaining items + Iterate::FlagHeads(linear_tid, head_flags, input, preds, flag_op); + + // Set tail_flags for remaining items + Iterate::FlagTails(linear_tid, tail_flags, input, flag_op); + } + + //! @rst + //! Sets both head and tail flags indicating discontinuities between items partitioned across the thread block. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The flag ``head_flags[i]`` is set for item ``input[i]`` when ``flag_op(previous-item, input[i])`` + //! returns ``true`` (where ``previous-item`` is either the preceding item in the same thread or the last item in + //! the previous thread). + //! - For *thread*\ :sub:`0`, item ``input[0]`` is compared against ``tile_predecessor_item``. + //! - The flag ``tail_flags[i]`` is set for item ``input[i]`` when ``flag_op(input[i], next-item)`` + //! returns ``true`` (where ``next-item`` is either the next item in the same thread or the first item in + //! the next thread). + //! - For *thread*\ :sub:`BLOCK_THREADS - 1`, item ``input[ITEMS_PER_THREAD - 1]`` is compared + //! against ``tile_successor_item``. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the head- and tail-flagging of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockDiscontinuity for a 1D block of 128 threads of type int + //! using BlockDiscontinuity = cub::BlockDiscontinuity; + //! + //! // Allocate shared memory for BlockDiscontinuity + //! __shared__ typename BlockDiscontinuity::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Have thread0 obtain the predecessor item for the entire tile + //! int tile_predecessor_item; + //! if (threadIdx.x == 0) tile_predecessor_item == ... + //! + //! // Have thread127 obtain the successor item for the entire tile + //! int tile_successor_item; + //! if (threadIdx.x == 127) tile_successor_item == ... + //! + //! // Collectively compute head and flags for discontinuities in the segment + //! int head_flags[4]; + //! int tail_flags[4]; + //! BlockDiscontinuity(temp_storage).FlagHeadsAndTails(head_flags, tile_predecessor_item, + //! tail_flags, tile_successor_item, + //! thread_data, cub::Inequality()); + //! } + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,0,1,1], [1,1,1,1], [2,3,3,3], ..., [124,125,125,125] }``, + //! that the ``tile_predecessor_item`` is ``0``, and that the + //! ``tile_successor_item`` is ``125``. The corresponding output ``head_flags`` + //! in those threads will be ``{ [0,0,1,0], [0,0,0,0], [1,1,0,0], [0,1,0,0], ... }``. + //! and the corresponding output ``tail_flags`` in those threads will be + //! ``{ [0,1,0,0], [0,0,0,1], [1,0,0,...], ..., [1,0,0,0] }``. + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam FlagT + //! **[inferred]** The flag type (must be an integer type) + //! + //! @tparam FlagOp + //! **[inferred]** Binary predicate functor type having member + //! `T operator()(const T &a, const T &b)` or member + //! `T operator()(const T &a, const T &b, unsigned int b_index)`, and returning `true` + //! if a discontinuity exists between `a` and `b`, otherwise `false`. `b_index` is the rank + //! of `b` in the aggregate tile of data. + //! + //! @param[out] head_flags + //! Calling thread's discontinuity head_flags + //! + //! @param[in] tile_predecessor_item + //! @rst + //! *thread*\ :sub:`0` only item with which to compare the first tile item (``input[0]`` from *thread*\ :sub:`0`). + //! @endrst + //! + //! @param[out] tail_flags + //! Calling thread's discontinuity tail_flags + //! + //! @param[in] tile_successor_item + //! @rst + //! *thread*\ :sub:`BLOCK_THREADS - 1` only item with which to compare the last tile item + //! (``input[ITEMS_PER_THREAD - 1]`` from *thread*\ :sub:`BLOCK_THREADS - 1`). + //! @endrst + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[in] flag_op + //! Binary boolean flag predicate + template + _CCCL_DEVICE _CCCL_FORCEINLINE void FlagHeadsAndTails( + FlagT (&head_flags)[ITEMS_PER_THREAD], + T tile_predecessor_item, + FlagT (&tail_flags)[ITEMS_PER_THREAD], + T tile_successor_item, + T (&input)[ITEMS_PER_THREAD], + FlagOp flag_op) + { + // Share first and last items + temp_storage.first_items[linear_tid] = input[0]; + temp_storage.last_items[linear_tid] = input[ITEMS_PER_THREAD - 1]; + + __syncthreads(); + + T preds[ITEMS_PER_THREAD]; + + // Set flag for first thread-item + preds[0] = (linear_tid == 0) ? tile_predecessor_item : // First thread + temp_storage.last_items[linear_tid - 1]; + + head_flags[0] = ApplyOp::FlagT(flag_op, preds[0], input[0], linear_tid * ITEMS_PER_THREAD); + + // Set flag for last thread-item + T successor_item = (linear_tid == BLOCK_THREADS - 1) ? tile_successor_item : // Last thread + temp_storage.first_items[linear_tid + 1]; + + tail_flags[ITEMS_PER_THREAD - 1] = ApplyOp::FlagT( + flag_op, input[ITEMS_PER_THREAD - 1], successor_item, (linear_tid * ITEMS_PER_THREAD) + ITEMS_PER_THREAD); + + // Set head_flags for remaining items + Iterate::FlagHeads(linear_tid, head_flags, input, preds, flag_op); + + // Set tail_flags for remaining items + Iterate::FlagTails(linear_tid, tail_flags, input, flag_op); + } + + //! @} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_exchange.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_exchange.cuh new file mode 100644 index 00000000..63e138e6 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_exchange.cuh @@ -0,0 +1,1314 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! The cub::BlockExchange class provides :ref:`collective ` methods for +//! rearranging data partitioned across a CUDA thread block. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! The BlockExchange class provides :ref:`collective ` methods for rearranging data partitioned +//! across a CUDA thread block. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - It is commonplace for blocks of threads to rearrange data items between threads. For example, the +//! device-accessible memory subsystem prefers access patterns where data items are "striped" across threads (where +//! consecutive threads access consecutive items), yet most block-wide operations prefer a "blocked" partitioning of +//! items across threads (where consecutive items belong to a single thread). +//! - BlockExchange supports the following types of data exchanges: +//! +//! - Transposing between :ref:`blocked ` and :ref:`striped ` +//! arrangements +//! - Transposing between :ref:`blocked ` and +//! :ref:`warp-striped ` arrangements +//! - Scattering ranked items to a :ref:`blocked arrangement ` +//! - Scattering ranked items to a :ref:`striped arrangement ` +//! +//! - @rowmajor +//! +//! A Simple Example +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @blockcollective{BlockExchange} +//! +//! The code snippet below illustrates the conversion from a "blocked" to a "striped" arrangement of 512 integer items +//! partitioned across 128 threads where each thread owns 4 items. +//! +//! .. code-block:: c++ +//! +//! #include // or equivalently +//! +//! __global__ void ExampleKernel(int *d_data, ...) +//! { +//! // Specialize BlockExchange for a 1D block of 128 threads owning 4 integer items each +//! using BlockExchange = cub::BlockExchange; +//! +//! // Allocate shared memory for BlockExchange +//! __shared__ typename BlockExchange::TempStorage temp_storage; +//! +//! // Load a tile of data striped across threads +//! int thread_data[4]; +//! cub::LoadDirectStriped<128>(threadIdx.x, d_data, thread_data); +//! +//! // Collectively exchange data into a blocked arrangement across threads +//! BlockExchange(temp_storage).StripedToBlocked(thread_data); +//! } +//! +//! Suppose the set of striped input ``thread_data`` across the block of threads is ``{ [0,128,256,384], +//! [1,129,257,385], ..., [127,255,383,511] }``. The corresponding output ``thread_data`` in those threads will be +//! ``{ [0,1,2,3], [4,5,6,7], [8,9,10,11], ..., [508,509,510,511] }``. +//! +//! Performance Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - Proper device-specific padding ensures zero bank conflicts for most types. +//! +//! Re-using dynamically allocating shared memory +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of dynamically shared memory with +//! BlockReduce and how to re-purpose the same memory region. This example can be easily adapted to the storage required +//! by BlockExchange. +//! @endrst +//! +//! @tparam T +//! The data type to be exchanged +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam ItemsPerThread +//! The number of items partitioned onto each thread. +//! +//! @tparam WarpTimeSlicing +//! **[optional]** When `true`, only use enough shared memory for a single warp's worth of +//! tile data, time-slicing the block-wide exchange over multiple synchronized rounds. Yields a smaller memory footprint +//! at the expense of decreased parallelism. (Default: false) +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockExchange +{ + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; ///< The thread block size in threads + static constexpr int WARP_THREADS = detail::warp_threads; + static constexpr int WARPS = (BLOCK_THREADS + WARP_THREADS - 1) / WARP_THREADS; // TODO(bgruber): use ceil_div in + // C++14 + static constexpr int LOG_SMEM_BANKS = detail::log2_smem_banks; + + static constexpr int TILE_ITEMS = BLOCK_THREADS * ItemsPerThread; + static constexpr int TIME_SLICES = WarpTimeSlicing ? WARPS : 1; + static constexpr int TIME_SLICED_THREADS = + WarpTimeSlicing ? ::cuda::std::min(BLOCK_THREADS, WARP_THREADS) : BLOCK_THREADS; + static constexpr int TIME_SLICED_ITEMS = TIME_SLICED_THREADS * ItemsPerThread; + static constexpr int WARP_TIME_SLICED_THREADS = ::cuda::std::min(BLOCK_THREADS, WARP_THREADS); + static constexpr int WARP_TIME_SLICED_ITEMS = WARP_TIME_SLICED_THREADS * ItemsPerThread; + + // Insert padding to avoid bank conflicts during raking when items per thread is a power of two and > 4 (otherwise + // we can typically use 128b loads) + static constexpr bool INSERT_PADDING = ItemsPerThread > 4 && ::cuda::is_power_of_two(ItemsPerThread); + static constexpr int PADDING_ITEMS = INSERT_PADDING ? (TIME_SLICED_ITEMS >> LOG_SMEM_BANKS) : 0; + + /// Shared memory storage layout type + struct alignas(16) _TempStorage + { + T buff[TIME_SLICED_ITEMS + PADDING_ITEMS]; + }; + +public: + /// @smemstorage{BlockExchange} + using TempStorage = Uninitialized<_TempStorage>; + +private: + _TempStorage& temp_storage; + + // TODO(bgruber): can we use signed int here? Only these variables are unsigned: + unsigned int linear_tid = RowMajorTid(BlockDimX, BlockDimY, BlockDimZ); + unsigned int lane_id = ::cuda::ptx::get_sreg_laneid(); + unsigned int warp_id = WARPS == 1 ? 0 : linear_tid / WARP_THREADS; + unsigned int warp_offset = warp_id * WARP_TIME_SLICED_ITEMS; + + /// Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + //! @brief Transposes data items from **blocked** arrangement to **striped** arrangement. Specialized for no + //! timeslicing. + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void BlockedToStriped( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + ::cuda::std::false_type /*time_slicing*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = linear_tid * ItemsPerThread + i; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i * BLOCK_THREADS + linear_tid; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + + //! @brief Transposes data items from **blocked** arrangement to **striped** arrangement. Specialized for + //! warp-timeslicing. + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void BlockedToStriped( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + ::cuda::std::true_type /*time_slicing*/) + { + T temp_items[ItemsPerThread]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int slice = 0; slice < TIME_SLICES; slice++) + { + const int slice_offset = slice * TIME_SLICED_ITEMS; + const int slice_oob = slice_offset + TIME_SLICED_ITEMS; + + __syncthreads(); + + if (warp_id == slice) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = lane_id * ItemsPerThread + i; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + // Read a strip of items + const int strip_offset = i * BLOCK_THREADS; + const int strip_oob = strip_offset + BLOCK_THREADS; + + if (slice_offset < strip_oob && slice_oob > strip_offset) + { + int item_offset = strip_offset + linear_tid - slice_offset; + if (item_offset >= 0 && item_offset < TIME_SLICED_ITEMS) + { + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + temp_items[i] = temp_storage.buff[item_offset]; + } + } + } + } + + // Copy + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + output_items[i] = temp_items[i]; + } + } + + //! @brief Transposes data items from **blocked** arrangement to **warp-striped** arrangement. Specialized for no + //! timeslicing + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void BlockedToWarpStriped( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + ::cuda::std::false_type /*time_slicing*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = warp_offset + i + (lane_id * ItemsPerThread); + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncwarp(0xffffffff); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = warp_offset + (i * WARP_TIME_SLICED_THREADS) + lane_id; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + + //! @brief Transposes data items from **blocked** arrangement to **warp-striped** arrangement. Specialized for + //! warp-timeslicing + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void BlockedToWarpStriped( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + ::cuda::std::true_type /*time_slicing*/) + { + if (warp_id == 0) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i + lane_id * ItemsPerThread; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncwarp(0xffffffff); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i * WARP_TIME_SLICED_THREADS + lane_id; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + + _CCCL_PRAGMA_UNROLL_FULL() + for (int slice = 1; slice < TIME_SLICES; ++slice) + { + __syncthreads(); + + if (warp_id == slice) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i + lane_id * ItemsPerThread; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncwarp(0xffffffff); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i * WARP_TIME_SLICED_THREADS + lane_id; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + } + } + + //! @brief Transposes data items from **striped** arrangement to **blocked** arrangement. Specialized for no + //! timeslicing. + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void StripedToBlocked( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + ::cuda::std::false_type /*time_slicing*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i * BLOCK_THREADS + linear_tid; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncthreads(); + + // No timeslicing + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = linear_tid * ItemsPerThread + i; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + + //! @brief Transposes data items from **striped** arrangement to **blocked** arrangement. Specialized for + //! warp-timeslicing. + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void StripedToBlocked( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + ::cuda::std::true_type /*time_slicing*/) + { + // Warp time-slicing + T temp_items[ItemsPerThread]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int slice = 0; slice < TIME_SLICES; slice++) + { + const int slice_offset = slice * TIME_SLICED_ITEMS; + const int slice_oob = slice_offset + TIME_SLICED_ITEMS; + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + // Write a strip of items + const int strip_offset = i * BLOCK_THREADS; + const int strip_oob = strip_offset + BLOCK_THREADS; + + if (slice_offset < strip_oob && slice_oob > strip_offset) + { + int item_offset = strip_offset + linear_tid - slice_offset; + if (item_offset >= 0 && item_offset < TIME_SLICED_ITEMS) + { + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + } + } + + __syncthreads(); + + if (warp_id == slice) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = lane_id * ItemsPerThread + i; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + temp_items[i] = temp_storage.buff[item_offset]; + } + } + } + + // Copy + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + output_items[i] = temp_items[i]; + } + } + + //! @brief Transposes data items from **warp-striped** arrangement to **blocked** arrangement. Specialized for no + //! timeslicing + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void WarpStripedToBlocked( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + ::cuda::std::false_type /*time_slicing*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = warp_offset + (i * WARP_TIME_SLICED_THREADS) + lane_id; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncwarp(0xffffffff); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = warp_offset + i + (lane_id * ItemsPerThread); + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(output_items + i, temp_storage.buff[item_offset]); + } + } + + //! @brief Transposes data items from **warp-striped** arrangement to **blocked** arrangement. Specialized for + //! warp-timeslicing + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void WarpStripedToBlocked( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + ::cuda::std::true_type /*time_slicing*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int slice = 0; slice < TIME_SLICES; ++slice) + { + __syncthreads(); + + if (warp_id == slice) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i * WARP_TIME_SLICED_THREADS + lane_id; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncwarp(0xffffffff); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i + lane_id * ItemsPerThread; + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + } + } + + //! @brief Exchanges data items annotated by rank into **blocked** arrangement. Specialized for no timeslicing. + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[in] ranks + //! Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToBlocked( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + OffsetT (&ranks)[ItemsPerThread], + ::cuda::std::false_type /*time_slicing*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = ranks[i]; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = linear_tid * ItemsPerThread + i; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + + //! @brief Exchanges data items annotated by rank into **blocked** arrangement. Specialized for warp-timeslicing. + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[in] ranks + //! Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToBlocked( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + OffsetT ranks[ItemsPerThread], + ::cuda::std::true_type /*time_slicing*/) + { + T temp_items[ItemsPerThread]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int slice = 0; slice < TIME_SLICES; slice++) + { + __syncthreads(); + + const int slice_offset = TIME_SLICED_ITEMS * slice; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = ranks[i] - slice_offset; + if (item_offset >= 0 && item_offset < WARP_TIME_SLICED_ITEMS) + { + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + } + + __syncthreads(); + + if (warp_id == slice) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = lane_id * ItemsPerThread + i; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + temp_items[i] = temp_storage.buff[item_offset]; + } + } + } + + // Copy + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + output_items[i] = temp_items[i]; + } + } + + //! @brief Exchanges data items annotated by rank into **striped** arrangement. Specialized for no timeslicing. + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[in] ranks + //! Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToStriped( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + OffsetT (&ranks)[ItemsPerThread], + ::cuda::std::false_type /*time_slicing*/) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = ranks[i]; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i * BLOCK_THREADS + linear_tid; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + + //! @brief Exchanges data items annotated by rank into **striped** arrangement. Specialized for warp-timeslicing. + //! + //! @param[in] input_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[out] output_items + //! Items to exchange, converting between **blocked** and **striped** arrangements. + //! + //! @param[in] ranks + //! Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToStriped( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + OffsetT (&ranks)[ItemsPerThread], + ::cuda::std::true_type /*time_slicing*/) + { + T temp_items[ItemsPerThread]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int slice = 0; slice < TIME_SLICES; slice++) + { + const int slice_offset = slice * TIME_SLICED_ITEMS; + const int slice_oob = slice_offset + TIME_SLICED_ITEMS; + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = ranks[i] - slice_offset; + if (item_offset >= 0 && item_offset < WARP_TIME_SLICED_ITEMS) + { + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + detail::uninitialized_copy_single(temp_storage.buff + item_offset, input_items[i]); + } + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + // Read a strip of items + const int strip_offset = i * BLOCK_THREADS; + const int strip_oob = strip_offset + BLOCK_THREADS; + + if (slice_offset < strip_oob && slice_oob > strip_offset) + { + int item_offset = strip_offset + linear_tid - slice_offset; + if (item_offset >= 0 && item_offset < TIME_SLICED_ITEMS) + { + if constexpr (INSERT_PADDING) + { + item_offset += item_offset >> LOG_SMEM_BANKS; + } + temp_items[i] = temp_storage.buff[item_offset]; + } + } + } + } + + // Copy + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + output_items[i] = temp_items[i]; + } + } + +public: + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using a private static allocation of shared memory as temporary storage. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + _CCCL_DEVICE _CCCL_FORCEINLINE BlockExchange() + : temp_storage(PrivateStorage()) + {} + + //! @brief Collective constructor using the specified memory allocation as temporary storage. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + //! + //! @param[in] temp_storage Reference to memory allocation having layout type TempStorage + _CCCL_DEVICE _CCCL_FORCEINLINE BlockExchange(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + {} + + //! @} + //! @name Structured exchanges + //! @{ + + //! @rst + //! Transposes data items from **striped** arrangement to **blocked** arrangement. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the conversion from a "striped" to a "blocked" arrangement + //! of 512 integer items partitioned across 128 threads where each thread owns 4 items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, ...) + //! { + //! // Specialize BlockExchange for a 1D block of 128 threads owning 4 integer items each + //! using BlockExchange = cub::BlockExchange; + //! + //! // Allocate shared memory for BlockExchange + //! __shared__ typename BlockExchange::TempStorage temp_storage; + //! + //! // Load a tile of ordered data into a striped arrangement across block threads + //! int thread_data[4]; + //! cub::LoadDirectStriped<128>(threadIdx.x, d_data, thread_data); + //! + //! // Collectively exchange data into a blocked arrangement across threads + //! BlockExchange(temp_storage).StripedToBlocked(thread_data, thread_data); + //! } + //! + //! Suppose the set of striped input ``thread_data`` across the block of threads is ``{ [0,128,256,384], + //! [1,129,257,385], ..., [127,255,383,511] }`` after loading from device-accessible memory. The corresponding output + //! ``thread_data`` in those threads will be ``{ [0,1,2,3], [4,5,6,7], [8,9,10,11], ..., [508,509,510,511] }``. + //! @endrst + //! + //! @param[in] input_items + //! Items to exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[out] output_items + //! Items from exchange, converting between **striped** and **blocked** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + StripedToBlocked(const T (&input_items)[ItemsPerThread], OutputT (&output_items)[ItemsPerThread]) + { + StripedToBlocked(input_items, output_items, detail::bool_constant_v); + } + + //! @rst + //! Transposes data items from **blocked** arrangement to **striped** arrangement. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the conversion from a "blocked" to a "striped" arrangement + //! of 512 integer items partitioned across 128 threads where each thread owns 4 items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, ...) + //! { + //! // Specialize BlockExchange for a 1D block of 128 threads owning 4 integer items each + //! using BlockExchange = cub::BlockExchange; + //! + //! // Allocate shared memory for BlockExchange + //! __shared__ typename BlockExchange::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively exchange data into a striped arrangement across threads + //! BlockExchange(temp_storage).BlockedToStriped(thread_data, thread_data); + //! + //! // Store data striped across block threads into an ordered tile + //! cub::StoreDirectStriped(threadIdx.x, d_data, thread_data); + //! } + //! + //! Suppose the set of blocked input ``thread_data`` across the block of threads is ``{ [0,1,2,3], [4,5,6,7], + //! [8,9,10,11], ..., [508,509,510,511] }``. The corresponding output ``thread_data`` in those threads will be + //! ``{ [0,128,256,384], [1,129,257,385], ..., [127,255,383,511] }`` in preparation for storing to device-accessible + //! memory. + //! @endrst + //! + //! @param[in] input_items + //! Items to exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[out] output_items + //! Items from exchange, converting between **striped** and **blocked** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + BlockedToStriped(const T (&input_items)[ItemsPerThread], OutputT (&output_items)[ItemsPerThread]) + { + BlockedToStriped(input_items, output_items, detail::bool_constant_v); + } + + //! @rst + //! Transposes data items from **warp-striped** arrangement to **blocked** arrangement. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @smemreuse + //! + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the conversion from a "warp-striped" to a "blocked" + //! arrangement of 512 integer items partitioned across 128 threads where each thread owns 4 + //! items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, ...) + //! { + //! // Specialize BlockExchange for a 1D block of 128 threads owning 4 integer items each + //! using BlockExchange = cub::BlockExchange; + //! + //! // Allocate shared memory for BlockExchange + //! __shared__ typename BlockExchange::TempStorage temp_storage; + //! + //! // Load a tile of ordered data into a warp-striped arrangement across warp threads + //! int thread_data[4]; + //! cub::LoadSWarptriped(threadIdx.x, d_data, thread_data); + //! + //! // Collectively exchange data into a blocked arrangement across threads + //! BlockExchange(temp_storage).WarpStripedToBlocked(thread_data); + //! } + //! + //! Suppose the set of warp-striped input ``thread_data`` across the block of threads is ``{ [0,32,64,96], + //! [1,33,65,97], [2,34,66,98], ..., [415,447,479,511] }`` after loading from device-accessible memory. (The first 128 + //! items are striped across the first warp of 32 threads, the second 128 items are striped across the second warp, + //! etc.) The corresponding output ``thread_data`` in those threads will be ``{ [0,1,2,3], [4,5,6,7], [8,9,10,11], + //! ..., [508,509,510,511] }``. + //! @endrst + //! + //! @param[in] input_items + //! Items to exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[out] output_items + //! Items from exchange, converting between **striped** and **blocked** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + WarpStripedToBlocked(const T (&input_items)[ItemsPerThread], OutputT (&output_items)[ItemsPerThread]) + { + WarpStripedToBlocked(input_items, output_items, detail::bool_constant_v); + } + + //! @rst + //! Transposes data items from **blocked** arrangement to **warp-striped** arrangement. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @smemreuse + //! + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the conversion from a "blocked" to a "warp-striped" + //! arrangement of 512 integer items partitioned across 128 threads where each thread owns 4 + //! items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, ...) + //! { + //! // Specialize BlockExchange for a 1D block of 128 threads owning 4 integer items each + //! using BlockExchange = cub::BlockExchange; + //! + //! // Allocate shared memory for BlockExchange + //! __shared__ typename BlockExchange::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively exchange data into a warp-striped arrangement across threads + //! BlockExchange(temp_storage).BlockedToWarpStriped(thread_data, thread_data); + //! + //! // Store data striped across warp threads into an ordered tile + //! cub::StoreDirectStriped(threadIdx.x, d_data, thread_data); + //! } + //! + //! Suppose the set of blocked input ``thread_data`` across the block of threads is ``{ [0,1,2,3], [4,5,6,7], + //! [8,9,10,11], ..., [508,509,510,511] }``. The corresponding output ``thread_data`` in those threads will be + //! ``{ [0,32,64,96], [1,33,65,97], [2,34,66,98], ..., [415,447,479,511] }`` in preparation for storing to + //! device-accessible memory. (The first 128 items are striped across the first warp of 32 threads, the second 128 + //! items are striped across the second warp, etc.) + //! @endrst + //! + //! @param[in] input_items + //! Items to exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[out] output_items + //! Items from exchange, converting between **striped** and **blocked** arrangements. + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + BlockedToWarpStriped(const T (&input_items)[ItemsPerThread], OutputT (&output_items)[ItemsPerThread]) + { + BlockedToWarpStriped(input_items, output_items, detail::bool_constant_v); + } + + //! @} + //! @name Scatter exchanges + //! @{ + + //! @rst + //! Exchanges data items annotated by rank into **blocked** arrangement. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @smemreuse + //! @endrst + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for local offsets + //! + //! @param[in] input_items + //! Items to exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[out] output_items + //! Items from exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[in] ranks + //! Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToBlocked( + const T (&input_items)[ItemsPerThread], OutputT (&output_items)[ItemsPerThread], OffsetT (&ranks)[ItemsPerThread]) + { + ScatterToBlocked(input_items, output_items, ranks, detail::bool_constant_v); + } + + //! @rst + //! Exchanges data items annotated by rank into **striped** arrangement. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @smemreuse + //! + //! @endrst + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for local offsets + //! + //! @param[in] input_items + //! Items to exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[out] output_items + //! Items from exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[in] ranks + //! Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToStriped( + const T (&input_items)[ItemsPerThread], OutputT (&output_items)[ItemsPerThread], OffsetT (&ranks)[ItemsPerThread]) + { + ScatterToStriped(input_items, output_items, ranks, detail::bool_constant_v); + } + + //! @rst + //! Exchanges data items annotated by rank into **striped** arrangement. Items with rank -1 are not exchanged. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @smemreuse + //! + //! @endrst + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for local offsets + //! + //! @param[in] input_items + //! Items to exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[out] output_items + //! Items from exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[in] ranks + //! Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToStripedGuarded( + const T (&input_items)[ItemsPerThread], OutputT (&output_items)[ItemsPerThread], OffsetT (&ranks)[ItemsPerThread]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = ranks[i]; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + if (ranks[i] >= 0) + { + temp_storage.buff[item_offset] = input_items[i]; + } + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i * BLOCK_THREADS + linear_tid; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + + //! @rst + //! Exchanges valid data items annotated by rank into **striped** arrangement. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @smemreuse + //! + //! @endrst + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for local offsets + //! + //! @tparam ValidFlag + //! **[inferred]** FlagT type denoting which items are valid + //! + //! @param[in] input_items + //! Items to exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[out] output_items + //! Items from exchange, converting between **striped** and **blocked** arrangements. + //! + //! @param[in] ranks + //! Corresponding scatter ranks + //! + //! @param[in] is_valid + //! Corresponding flag denoting item validity + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToStripedFlagged( + const T (&input_items)[ItemsPerThread], + OutputT (&output_items)[ItemsPerThread], + OffsetT (&ranks)[ItemsPerThread], + ValidFlag (&is_valid)[ItemsPerThread]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = ranks[i]; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + if (is_valid[i]) + { + temp_storage.buff[item_offset] = input_items[i]; + } + } + + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + int item_offset = i * BLOCK_THREADS + linear_tid; + if constexpr (INSERT_PADDING) + { + item_offset = (item_offset >> LOG_SMEM_BANKS) + item_offset; + } + output_items[i] = temp_storage.buff[item_offset]; + } + } + + //! @} + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + + /// @param[in-out] items + /// Items to exchange, converting between **striped** and **blocked** arrangements. + _CCCL_DEVICE _CCCL_FORCEINLINE void StripedToBlocked(T (&items)[ItemsPerThread]) + { + StripedToBlocked(items, items); + } + + /// @param[in-out] items + /// Items to exchange, converting between **striped** and **blocked** arrangements. + _CCCL_DEVICE _CCCL_FORCEINLINE void BlockedToStriped(T (&items)[ItemsPerThread]) + { + BlockedToStriped(items, items); + } + + /// @param[in-out] items + /// Items to exchange, converting between **striped** and **blocked** arrangements. + _CCCL_DEVICE _CCCL_FORCEINLINE void WarpStripedToBlocked(T (&items)[ItemsPerThread]) + { + WarpStripedToBlocked(items, items); + } + + /// @param[in-out] items + /// Items to exchange, converting between **striped** and **blocked** arrangements. + _CCCL_DEVICE _CCCL_FORCEINLINE void BlockedToWarpStriped(T (&items)[ItemsPerThread]) + { + BlockedToWarpStriped(items, items); + } + + /// @param[in-out] items + /// Items to exchange, converting between **striped** and **blocked** arrangements. + /// + /// @param[in] ranks + /// Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToBlocked(T (&items)[ItemsPerThread], OffsetT (&ranks)[ItemsPerThread]) + { + ScatterToBlocked(items, items, ranks); + } + + /// @param[in-out] items + /// Items to exchange, converting between **striped** and **blocked** arrangements. + /// @param[in] ranks + /// Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToStriped(T (&items)[ItemsPerThread], OffsetT (&ranks)[ItemsPerThread]) + { + ScatterToStriped(items, items, ranks); + } + + /// @param[in-out] items + /// Items to exchange, converting between **striped** and **blocked** arrangements. + /// @param[in] ranks + /// Corresponding scatter ranks + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ScatterToStripedGuarded(T (&items)[ItemsPerThread], OffsetT (&ranks)[ItemsPerThread]) + { + ScatterToStripedGuarded(items, items, ranks); + } + + /// @param[in-out] items + /// Items to exchange, converting between **striped** and **blocked** arrangements. + /// @param[in] ranks + /// Corresponding scatter ranks + /// @param[in] is_valid + /// Corresponding flag denoting item validity + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ScatterToStripedFlagged( + T (&items)[ItemsPerThread], OffsetT (&ranks)[ItemsPerThread], ValidFlag (&is_valid)[ItemsPerThread]) + { + ScatterToStripedFlagged(items, items, ranks, is_valid); + } + +#endif // _CCCL_DOXYGEN_INVOKED +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_histogram.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_histogram.cuh new file mode 100644 index 00000000..f2eef8d1 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_histogram.cuh @@ -0,0 +1,412 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * The cub::BlockHistogram class provides [collective](../index.html#sec0) methods for + * constructing block-wide histograms from data samples partitioned across a CUDA thread block. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +//! @brief BlockHistogramAlgorithm enumerates alternative algorithms for the parallel construction of +//! block-wide histograms. +enum BlockHistogramAlgorithm +{ + + //! @rst + //! + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! Sorting followed by differentiation. Execution is comprised of two phases: + //! + //! #. Sort the data using efficient radix sort + //! #. Look for "runs" of same-valued keys by detecting discontinuities; the run-lengths are histogram bin counts. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! Delivers consistent throughput regardless of sample bin distribution. + //! + //! @endrst + BLOCK_HISTO_SORT, + + //! @rst + //! + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! Use atomic addition to update byte counts directly + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! Performance is strongly tied to the hardware implementation of atomic + //! addition, and may be significantly degraded for non uniformly-random + //! input distributions where many concurrent updates are likely to be + //! made to the same bin counter. + //! + //! @endrst + BLOCK_HISTO_ATOMIC, +}; + +//! @rst +//! The BlockHistogram class provides :ref:`collective ` methods for +//! constructing block-wide histograms from data samples partitioned across a CUDA thread block. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - A `histogram `_ counts the number of observations that fall into +//! each of the disjoint categories (known as *bins*). +//! - The ``T`` type must be implicitly castable to an integer type. +//! - BlockHistogram expects each integral ``input[i]`` value to satisfy +//! ``0 <= input[i] < Bins``. Values outside of this range result in undefined behavior. +//! - BlockHistogram can be optionally specialized to use different algorithms: +//! +//! #. :cpp:enumerator:`cub::BLOCK_HISTO_SORT`: Sorting followed by differentiation. +//! #. :cpp:enumerator:`cub::BLOCK_HISTO_ATOMIC`: Use atomic addition to update byte counts directly. +//! +//! A Simple Example +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @blockcollective{BlockHistogram} +//! +//! The code snippet below illustrates a 256-bin histogram of 512 integer samples that +//! are partitioned across 128 threads where each thread owns 4 samples. +//! +//! .. code-block:: c++ +//! +//! #include // or equivalently +//! +//! __global__ void ExampleKernel(...) +//! { +//! // Specialize a 256-bin BlockHistogram type for a 1D block of 128 threads having 4 character samples each +//! using BlockHistogram = cub::BlockHistogram; +//! +//! // Allocate shared memory for BlockHistogram +//! __shared__ typename BlockHistogram::TempStorage temp_storage; +//! +//! // Allocate shared memory for block-wide histogram bin counts +//! __shared__ unsigned int smem_histogram[256]; +//! +//! // Obtain input samples per thread +//! unsigned char data[4]; +//! ... +//! +//! // Compute the block-wide histogram +//! BlockHistogram(temp_storage).Histogram(data, smem_histogram); +//! } +//! +//! Performance and Usage Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - @granularity +//! - All input values must fall between ``[0, Bins)``, or behavior is undefined. +//! - The histogram output can be constructed in shared or device-accessible memory +//! - See ``cub::BlockHistogramAlgorithm`` for performance details regarding algorithmic alternatives +//! +//! Re-using dynamically allocating shared memory +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of dynamically shared memory with +//! BlockReduce and how to re-purpose the same memory region. This example can be easily adapted to the storage +//! required by BlockHistogram. +//! @endrst +//! +//! @tparam T +//! The sample type being histogrammed (must be castable to an integer bin identifier) +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam ItemsPerThread +//! The number of items per thread +//! +//! @tparam Bins +//! The number bins within the histogram +//! +//! @tparam Algorithm +//! **[optional]** cub::BlockHistogramAlgorithm enumerator specifying the underlying algorithm to use +//! (default: cub::BLOCK_HISTO_SORT) +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockHistogram +{ +private: + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + /// Internal specialization. + using InternalBlockHistogram = + ::cuda::std::_If, + detail::BlockHistogramAtomic>; + + /// Shared memory storage layout type for BlockHistogram + using _TempStorage = typename InternalBlockHistogram::TempStorage; + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + unsigned int linear_tid; + + /// Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + +public: + /// @smemstorage{BlockHistogram} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using a private static allocation of shared memory as temporary storage. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + _CCCL_DEVICE _CCCL_FORCEINLINE BlockHistogram() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @brief Collective constructor using the specified memory allocation as temporary storage. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @param[in] temp_storage + * Reference to memory allocation having layout type TempStorage + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockHistogram(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Histogram operations + //! @{ + + //! @rst + //! Initialize the shared histogram counters to zero. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a the initialization and update of a + //! histogram of 512 integer samples that are partitioned across 128 threads + //! where each thread owns 4 samples. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize a 256-bin BlockHistogram type for a 1D block of 128 threads having 4 character samples each + //! using BlockHistogram = cub::BlockHistogram; + //! + //! // Allocate shared memory for BlockHistogram + //! __shared__ typename BlockHistogram::TempStorage temp_storage; + //! + //! // Allocate shared memory for block-wide histogram bin counts + //! __shared__ unsigned int smem_histogram[256]; + //! + //! // Obtain input samples per thread + //! unsigned char thread_samples[4]; + //! ... + //! + //! // Initialize the block-wide histogram + //! BlockHistogram(temp_storage).InitHistogram(smem_histogram); + //! + //! // Update the block-wide histogram + //! BlockHistogram(temp_storage).Composite(thread_samples, smem_histogram); + //! } + //! @endrst + //! + //! @tparam CounterT + //! **[inferred]** Histogram counter type + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InitHistogram(CounterT histogram[Bins]) + { + // Initialize histogram bin counts to zeros + int histo_offset = 0; + + _CCCL_PRAGMA_UNROLL_FULL() + for (; histo_offset + BLOCK_THREADS <= Bins; histo_offset += BLOCK_THREADS) + { + histogram[histo_offset + linear_tid] = 0; + } + // Finish up with guarded initialization if necessary + if ((Bins % BLOCK_THREADS != 0) && (histo_offset + linear_tid < Bins)) + { + histogram[histo_offset + linear_tid] = 0; + } + } + + //! @rst + //! Constructs a block-wide histogram in shared/device-accessible memory. + //! Each thread contributes an array of input elements. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a 256-bin histogram of 512 integer samples that + //! are partitioned across 128 threads where each thread owns 4 samples. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize a 256-bin BlockHistogram type for a 1D block of 128 threads having 4 character samples each + //! using BlockHistogram = cub::BlockHistogram; + //! + //! // Allocate shared memory for BlockHistogram + //! __shared__ typename BlockHistogram::TempStorage temp_storage; + //! + //! // Allocate shared memory for block-wide histogram bin counts + //! __shared__ unsigned int smem_histogram[256]; + //! + //! // Obtain input samples per thread + //! unsigned char thread_samples[4]; + //! ... + //! + //! // Compute the block-wide histogram + //! BlockHistogram(temp_storage).Histogram(thread_samples, smem_histogram); + //! } + //! @endrst + //! + //! @tparam CounterT + //! **[inferred]** Histogram counter type + //! + //! @param[in] items + //! Calling thread's input values to histogram + //! + //! @param[out] histogram + //! Reference to shared/device-accessible memory histogram + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Histogram(T (&items)[ItemsPerThread], CounterT histogram[Bins]) + { + // Initialize histogram bin counts to zeros + InitHistogram(histogram); + + __syncthreads(); + + // Composite the histogram + InternalBlockHistogram(temp_storage).Composite(items, histogram); + } + + //! @rst + //! Updates an existing block-wide histogram in shared/device-accessible memory. + //! Each thread composites an array of input elements. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a the initialization and update of a + //! histogram of 512 integer samples that are partitioned across 128 threads + //! where each thread owns 4 samples. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize a 256-bin BlockHistogram type for a 1D block of 128 threads having 4 character samples each + //! using BlockHistogram = cub::BlockHistogram; + //! + //! // Allocate shared memory for BlockHistogram + //! __shared__ typename BlockHistogram::TempStorage temp_storage; + //! + //! // Allocate shared memory for block-wide histogram bin counts + //! __shared__ unsigned int smem_histogram[256]; + //! + //! // Obtain input samples per thread + //! unsigned char thread_samples[4]; + //! ... + //! + //! // Initialize the block-wide histogram + //! BlockHistogram(temp_storage).InitHistogram(smem_histogram); + //! + //! // Update the block-wide histogram + //! BlockHistogram(temp_storage).Composite(thread_samples, smem_histogram); + //! } + //! @endrst + //! + //! @tparam CounterT + //! **[inferred]** Histogram counter type + //! + //! @param[in] items + //! Calling thread's input values to histogram + //! + //! @param[out] histogram + //! Reference to shared/device-accessible memory histogram + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Composite(T (&items)[ItemsPerThread], CounterT histogram[Bins]) + { + InternalBlockHistogram(temp_storage).Composite(items, histogram); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_load.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_load.cuh new file mode 100644 index 00000000..077abb17 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_load.cuh @@ -0,0 +1,1192 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2016, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! block_load.cuh Operations for reading linear tiles of data into the CUDA thread block. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @name Blocked arrangement I/O (direct) +//! @{ + +//! @rst +//! Load a linear segment of items into a blocked arrangement across the thread block. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @blocked +//! +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to load. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **[inferred]** The random-access iterator type for input iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D +//! thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base input iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +LoadDirectBlocked(int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread]) +{ + // Load directly in thread-blocked order + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + dst_items[i] = block_src_it[linear_tid * ItemsPerThread + i]; // NOLINT(bugprone-misplaced-widening-cast) + } +} + +//! @rst +//! Load a linear segment of items into a blocked arrangement across the thread block, guarded by range. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @blocked +//! +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to load. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **[inferred]** The random-access iterator type for input iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D +//! thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +//! +//! @param[in] block_items_end +//! First out-of-bounds index when loading from block_src_it +template +_CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectBlocked( + int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread], int block_items_end) +{ + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + const auto src_pos = linear_tid * ItemsPerThread + i; + if (src_pos < block_items_end) + { + dst_items[i] = block_src_it[src_pos]; + } + } +} + +//! @rst +//! Load a linear segment of items into a blocked arrangement across the thread block, guarded +//! by range, with a fall-back assignment of out-of-bound elements. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @blocked +//! +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to load. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **[inferred]** The random-access iterator type for input \iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D +//! thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base input iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +//! +//! @param[in] block_items_end +//! First out-of-bounds index when loading from block_src_it +//! +//! @param[in] oob_default +//! Default value to assign out-of-bound items +template +_CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectBlocked( + int linear_tid, + RandomAccessIterator block_src_it, + T (&dst_items)[ItemsPerThread], + int block_items_end, + DefaultT oob_default) +{ + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + dst_items[i] = oob_default; + } + + LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end); +} + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + +//! @brief Internal implementation for load vectorization +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D +//! thread blocks) +//! +//! @param[in] block_src_ptr +//! Input pointer for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (&dst_items)[ItemsPerThread]) +{ + // Find biggest memory access word that T is a whole multiple of + using device_word_t = typename UnitWord::DeviceWord; + _CCCL_DIAG_PUSH +# if _CCCL_COMPILER(CLANG, >=, 10) + _CCCL_DIAG_SUPPRESS_CLANG("-Wsizeof-array-div") +# endif // _CCCL_COMPILER(CLANG, >=, 10) + // NOLINTNEXTLINE(bugprone-sizeof-expression) + constexpr int total_words = static_cast(sizeof(dst_items) / sizeof(device_word_t)); + _CCCL_DIAG_POP + constexpr int vector_size = (total_words % 4 == 0) ? 4 : (total_words % 2 == 0) ? 2 : 1; + constexpr int vectors_per_thread = total_words / vector_size; + + // Load into an array of vectors in thread-blocked order + using vector_t = typename CubVector::Type; + + // Add the alignment check to ensure the vectorized loading can proceed. + if (::cuda::std::is_sufficiently_aligned(block_src_ptr)) + { + vector_t vec_items[vectors_per_thread]; + // Load into an array of vectors in thread-blocked order + const vector_t* vec_ptr = reinterpret_cast(block_src_ptr) + linear_tid * vectors_per_thread; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < vectors_per_thread; i++) + { + vec_items[i] = ThreadLoad(vec_ptr + i); + } + + // Copy to destination + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + dst_items[i] = *(reinterpret_cast(vec_items) + i); + } + } + else + { + LoadDirectBlocked(linear_tid, block_src_ptr, dst_items); + } +} + +#endif // _CCCL_DOXYGEN_INVOKED + +//! @rst +//! Load a linear segment of items into a blocked arrangement across the thread block. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @blocked +//! +//! The input offset (``block_ptr + block_offset``) must be quad-item aligned +//! +//! The following conditions will prevent vectorization and loading will fall back to cub::BLOCK_LOAD_DIRECT: +//! +//! - ``ItemsPerThread`` is odd +//! - The data type ``T`` is not a built-in primitive or CUDA vector type +//! (e.g., ``short``, ``int2``, ``double``, ``float2``, etc.) +//! +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to load. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + +//! linear_tid` for 2D thread blocks) +//! +//! @param[in] block_src_ptr +//! The thread block's base pointer for loading from +//! +//! @param[out] dst_items +//! destination to load data into +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +LoadDirectBlockedVectorized(int linear_tid, T* block_src_ptr, T (&dst_items)[ItemsPerThread]) +{ + InternalLoadDirectBlockedVectorized(linear_tid, block_src_ptr, dst_items); +} + +//! @} +//! @name Striped arrangement I/O (direct) +//! @{ + +//! @rst +//! Load a linear segment of items into a striped arrangement across the thread block. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @striped +//! +//! @endrst +//! +//! @tparam ThreadsPerBlock +//! The thread block size in threads +//! +//! @tparam T +//! **[inferred]** The data type to load. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **[inferred]** The random-access iterator type for input iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D +//! thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +LoadDirectStriped(int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread]) +{ + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + dst_items[i] = block_src_it[linear_tid + i * ThreadsPerBlock]; // NOLINT(bugprone-misplaced-widening-cast) + } +} + +namespace detail +{ +template +_CCCL_DEVICE _CCCL_FORCEINLINE void load_transform_direct_striped( + int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread], TransformOpT transform_op) +{ + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + dst_items[i] = + transform_op(block_src_it[linear_tid + i * ThreadsPerBlock]); // NOLINT(bugprone-misplaced-widening-cast) + } +} +} // namespace detail + +//! @rst +//! Load a linear segment of items into a striped arrangement across the thread block, guarded by range +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @striped +//! +//! @endrst +//! +//! @tparam ThreadsPerBlock +//! The thread block size in threads +//! +//! @tparam T +//! **inferred** The data type to load. +//! +//! @tparam ItemsPerThread +//! **inferred** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **inferred** The random-access iterator type for input iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., (threadIdx.y * blockDim.x) + +//! linear_tid for 2D thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +//! +//! @param[in] block_items_end +//! Number of valid items to load +template +_CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectStriped( + int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread], int block_items_end) +{ + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + const auto src_pos = linear_tid + i * ThreadsPerBlock; + if (src_pos < block_items_end) + { + dst_items[i] = block_src_it[src_pos]; + } + } +} + +//! @rst +//! Load a linear segment of items into a striped arrangement across the thread block, guarded +//! by range, with a fall-back assignment of out-of-bound elements. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @striped +//! +//! @endrst +//! +//! @tparam ThreadsPerBlock +//! The thread block size in threads +//! +//! @tparam T +//! **inferred** The data type to load. +//! +//! @tparam ItemsPerThread +//! **inferred** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **inferred** The random-access iterator type for input \iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + +//! linear_tid` for 2D thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +//! +//! @param[in] block_items_end +//! Number of valid items to load +//! +//! @param[in] oob_default +//! Default value to assign out-of-bound items +template +_CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectStriped( + int linear_tid, + RandomAccessIterator block_src_it, + T (&dst_items)[ItemsPerThread], + int block_items_end, + DefaultT oob_default) +{ + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + dst_items[i] = oob_default; + } + + LoadDirectStriped(linear_tid, block_src_it, dst_items, block_items_end); +} + +//! @} +//! @name Warp-striped arrangement I/O (direct) +//! @{ + +//! @rst +//! Load a linear segment of items into a warp-striped arrangement across the thread block. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @warpstriped +//! +//! Usage Considerations +//! ++++++++++++++++++++ +//! +//! The number of threads in the thread block must be a multiple of the architecture's warp size. +//! +//! @endrst +//! +//! @tparam T +//! **inferred** The data type to load. +//! +//! @tparam ItemsPerThread +//! **inferred** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **inferred** The random-access iterator type for input iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + +//! linear_tid` for 2D thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +LoadDirectWarpStriped(int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread]) +{ + const int tid = linear_tid & (detail::warp_threads - 1); + const int wid = linear_tid >> detail::log2_warp_threads; + const int warp_offset = wid * detail::warp_threads * ItemsPerThread; + + // Load directly in warp-striped order + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + new (&dst_items[i]) + T(block_src_it[warp_offset + tid + (i * detail::warp_threads)]); // NOLINT(bugprone-misplaced-widening-cast) + } +} + +//! @rst +//! Load a linear segment of items into a warp-striped arrangement across the thread block, guarded by range +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @warpstriped +//! +//! Usage Considerations +//! ++++++++++++++++++++ +//! +//! The number of threads in the thread block must be a multiple of the architecture's warp size. +//! +//! @endrst +//! +//! @tparam T +//! **inferred** The data type to load. +//! +//! @tparam ItemsPerThread +//! **inferred** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **inferred** The random-access iterator type for input \iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + +//! linear_tid` for 2D thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +//! +//! @param[in] block_items_end +//! Number of valid items to load +template +_CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectWarpStriped( + int linear_tid, RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread], int block_items_end) +{ + const int tid = linear_tid & (detail::warp_threads - 1); + const int wid = linear_tid >> detail::log2_warp_threads; + const int warp_offset = wid * detail::warp_threads * ItemsPerThread; + + // Load directly in warp-striped order + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + const auto src_pos = warp_offset + tid + (i * detail::warp_threads); + if (src_pos < block_items_end) + { + new (&dst_items[i]) T(block_src_it[src_pos]); + } + } +} + +//! @rst +//! Load a linear segment of items into a warp-striped arrangement across the thread block, +//! guarded by range, with a fall-back assignment of out-of-bound elements. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @warpstriped +//! +//! @endrst +//! +//! Usage Considerations +//! ++++++++++++++++++++ +//! +//! The number of threads in the thread block must be a multiple of the architecture's warp size. +//! +//! @tparam T +//! **inferred** The data type to load. +//! +//! @tparam ItemsPerThread +//! **inferred** The number of consecutive items partitioned onto each thread. +//! +//! @tparam RandomAccessIterator +//! **inferred** The random-access iterator type for input \iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + +//! linear_tid` for 2D thread blocks) +//! +//! @param[in] block_src_it +//! The thread block's base iterator for loading from +//! +//! @param[out] dst_items +//! Destination to load data into +//! +//! @param[in] block_items_end +//! Number of valid items to load +//! +//! @param[in] oob_default +//! Default value to assign out-of-bound items +template +_CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectWarpStriped( + int linear_tid, + RandomAccessIterator block_src_it, + T (&dst_items)[ItemsPerThread], + int block_items_end, + DefaultT oob_default) +{ + // Load directly in warp-striped order + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ItemsPerThread; i++) + { + dst_items[i] = oob_default; + } + + LoadDirectWarpStriped(linear_tid, block_src_it, dst_items, block_items_end); +} + +//! @} + +//! @brief cub::BlockLoadAlgorithm enumerates alternative algorithms for cub::BlockLoad to read a linear segment of data +//! from memory into a blocked arrangement across a CUDA thread block. +enum BlockLoadAlgorithm +{ + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`blocked arrangement ` of data is read directly from memory. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! The utilization of memory transactions (coalescing) decreases as the access stride between threads increases + //! (i.e., the number items per thread). + //! @endrst + BLOCK_LOAD_DIRECT, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`striped arrangement ` of data is read directly from memory. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! The utilization of memory transactions (coalescing) doesn't depend on the number of items per thread. + //! + //! @endrst + BLOCK_LOAD_STRIPED, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`blocked arrangement ` of data is read from memory using CUDA's built-in + //! vectorized loads as a coalescing optimization. For example, ``ld.global.v4.s32`` instructions will be generated + //! when ``T = int`` and ``ItemsPerThread % 4 == 0``. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) remains high until the the access stride between threads + //! (i.e., the number items per thread) exceeds the maximum vector load width (typically 4 items or 64B, whichever + //! is lower). + //! - The following conditions will prevent vectorization and loading will fall back to cub::BLOCK_LOAD_DIRECT: + //! + //! - ``ItemsPerThread`` is odd + //! - The ``RandomAccessIterator`` is not a simple pointer type + //! - The block input offset is not quadword-aligned + //! - The data type ``T`` is not a built-in primitive or CUDA vector type + //! (e.g., ``short``, ``int2``, ``double``, ``float2``, etc.) + //! + //! @endrst + BLOCK_LOAD_VECTORIZE, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`striped arrangement ` of data is read efficiently from memory and then locally + //! transposed into a :ref:`blocked arrangement `. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) remains high regardless of items loaded per thread. + //! - The local reordering incurs slightly longer latencies and throughput than the direct cub::BLOCK_LOAD_DIRECT and + //! cub::BLOCK_LOAD_VECTORIZE alternatives. + //! + //! @endrst + BLOCK_LOAD_TRANSPOSE, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`warp-striped arrangement ` of data is read efficiently from memory and then + //! locally transposed into a :ref:`blocked arrangement `. + //! + //! Usage Considerations + //! ++++++++++++++++++++++++++ + //! + //! - ThreadsPerBlock must be a multiple of WARP_THREADS + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) remains high regardless of items loaded per thread. + //! - The local reordering incurs slightly larger latencies than the direct cub::BLOCK_LOAD_DIRECT and + //! cub::BLOCK_LOAD_VECTORIZE alternatives. + //! - Provisions more shared storage, but incurs smaller latencies than the + //! BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED alternative. + //! + //! @endrst + BLOCK_LOAD_WARP_TRANSPOSE, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! Like ``BLOCK_LOAD_WARP_TRANSPOSE``, a :ref:`warp-striped arrangement ` of data is read + //! directly from memory and then is locally transposed into a :ref:`blocked arrangement `. + //! To reduce the shared memory requirement, only one warp's worth of shared memory is provisioned and is subsequently + //! time-sliced among warps. + //! + //! Usage Considerations + //! ++++++++++++++++++++++++++ + //! + //! - ThreadsPerBlock must be a multiple of WARP_THREADS + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) remains high regardless of items loaded per thread. + //! - Provisions less shared memory temporary storage, but incurs larger latencies than the BLOCK_LOAD_WARP_TRANSPOSE + //! alternative. + //! + //! @endrst + BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED, +}; + +#if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED) +namespace detail +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(BlockLoadAlgorithm algo) noexcept +{ + switch (algo) + { + case BLOCK_LOAD_DIRECT: + return "BLOCK_LOAD_DIRECT"; + case BLOCK_LOAD_STRIPED: + return "BLOCK_LOAD_STRIPED"; + case BLOCK_LOAD_VECTORIZE: + return "BLOCK_LOAD_VECTORIZE"; + case BLOCK_LOAD_TRANSPOSE: + return "BLOCK_LOAD_TRANSPOSE"; + case BLOCK_LOAD_WARP_TRANSPOSE: + return "BLOCK_LOAD_WARP_TRANSPOSE"; + case BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED: + return "BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED"; + } + return ""; +} +} // namespace detail + +inline ::std::ostream& operator<<(::std::ostream& os, BlockLoadAlgorithm algo) +{ + return os << CUB_NS_QUALIFIER::detail::to_string(algo); +} +#endif // _CCCL_HOSTED() && !_CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END + +#if __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) +template <::cuda::std::same_as CharT> +struct std::formatter : formatter +{ + template + auto format(const CUB_NS_QUALIFIER::BlockLoadAlgorithm& algo, FmtCtx& ctx) const + { + return formatter::format(CUB_NS_QUALIFIER::detail::to_string(algo), ctx); + } +}; +#endif // __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) + +CUB_NAMESPACE_BEGIN + +//! @rst +//! The BlockLoad class provides :ref:`collective ` data movement methods for loading a linear +//! segment of items from memory into a :ref:`blocked arrangement ` across a CUDA thread +//! block. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - The BlockLoad class provides a single data movement abstraction that can be specialized to implement different +//! cub::BlockLoadAlgorithm strategies. This facilitates different performance policies for different architectures, +//! data types, granularity sizes, etc. +//! - BlockLoad can be optionally specialized by different data movement strategies: +//! +//! #. :cpp:enumerator:`cub::BLOCK_LOAD_DIRECT`: +//! A :ref:`blocked arrangement ` of data is read directly from memory. +//! #. :cpp:enumerator:`cub::BLOCK_LOAD_STRIPED`: +//! A :ref:`striped arrangement ` of data is read directly from memory. +//! #. :cpp:enumerator:`cub::BLOCK_LOAD_VECTORIZE`: +//! A :ref:`blocked arrangement ` of data is read directly from memory +//! using CUDA's built-in vectorized loads as a coalescing optimization. +//! #. :cpp:enumerator:`cub::BLOCK_LOAD_TRANSPOSE`: +//! A :ref:`striped arrangement ` of data is read directly from memory and is then +//! locally transposed into a :ref:`blocked arrangement `. +//! #. :cpp:enumerator:`cub::BLOCK_LOAD_WARP_TRANSPOSE`: +//! A :ref:`warp-striped arrangement ` of data is read directly from memory and is then +//! locally transposed into a :ref:`blocked arrangement `. +//! #. :cpp:enumerator:`cub::BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED`: +//! A :ref:`warp-striped arrangement ` of data is read directly from memory and is then +//! locally transposed into a :ref:`blocked arrangement ` one warp at a time. +//! +//! - @rowmajor +//! +//! A Simple Example +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @blockcollective{BlockLoad} +//! +//! The code snippet below illustrates the loading of a linear segment of 512 integers into a "blocked" arrangement +//! across 128 threads where each thread owns 4 consecutive items. The load is specialized for +//! ``BLOCK_LOAD_WARP_TRANSPOSE``, meaning memory references are efficiently coalesced using a warp-striped access +//! pattern (after which items are locally reordered among threads). +//! +//! .. code-block:: c++ +//! +//! #include // or equivalently +//! +//! __global__ void ExampleKernel(int *d_data, ...) +//! { +//! // Specialize BlockLoad for a 1D block of 128 threads owning 4 integer items each +//! using BlockLoad = cub::BlockLoad; +//! +//! // Allocate shared memory for BlockLoad +//! __shared__ typename BlockLoad::TempStorage temp_storage; +//! +//! // Load a segment of consecutive items that are blocked across threads +//! int thread_data[4]; +//! BlockLoad(temp_storage).Load(d_data, thread_data); +//! } +//! +//! Suppose the input ``d_data`` is ``0, 1, 2, 3, 4, 5, ...``. The set of ``thread_data`` across the block of threads in +//! those threads will be ``{ [0,1,2,3], [4,5,6,7], ..., [508,509,510,511] }``. +//! +//! Re-using dynamically allocating shared memory +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of dynamically shared memory with +//! BlockReduce and how to re-purpose the same memory region. This example can be easily adapted to the storage required +//! by BlockLoad. +//! +//! @endrst +//! +//! @tparam T +//! The data type to read into (which must be convertible from the input iterator's value type). +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam ItemsPerThread +//! The number of consecutive items partitioned onto each thread. +//! +//! @tparam Algorithm +//! **[optional]** cub::BlockLoadAlgorithm tuning policy. default: ``cub::BLOCK_LOAD_DIRECT``. +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockLoad +{ + static constexpr int ThreadsPerBlock = BlockDimX * BlockDimY * BlockDimZ; // total threads in the block + + // transposing load algorithms need a BlockExchange + using block_exchange = + BlockExchange; + + static_assert((Algorithm != BLOCK_LOAD_WARP_TRANSPOSE && Algorithm != BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED) + || (ThreadsPerBlock % detail::warp_threads == 0), + "ThreadsPerBlock must be a multiple of warp_threads for this BlockLoadAlgorithm"); + + _CCCL_HOST_DEVICE_API static constexpr auto temp_storage_helper() + { + if constexpr (Algorithm == BLOCK_LOAD_DIRECT || Algorithm == BLOCK_LOAD_STRIPED + || Algorithm == BLOCK_LOAD_VECTORIZE) + { + return NullType{}; + } + else if constexpr (Algorithm == BLOCK_LOAD_TRANSPOSE || Algorithm == BLOCK_LOAD_WARP_TRANSPOSE + || Algorithm == BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED) + { + return typename block_exchange::TempStorage{}; + } + } + + using _TempStorage = decltype(temp_storage_helper()); + + // Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + _TempStorage& temp_storage; + int linear_tid; + +public: + /// @smemstorage{BlockLoad} + using TempStorage = Uninitialized<_TempStorage>; + + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using a private static allocation of shared memory as temporary storage. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + _CCCL_DEVICE _CCCL_FORCEINLINE BlockLoad() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @brief Collective constructor using the specified memory allocation as temporary storage. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + //! + //! @param[in] temp_storage Reference to memory allocation having layout type TempStorage + _CCCL_DEVICE _CCCL_FORCEINLINE BlockLoad(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Data movement + //! @{ + + //! @rst + //! Load a linear segment of items from memory. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @blocked + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the loading of a linear segment of 512 integers into a "blocked" arrangement + //! across 128 threads where each thread owns 4 consecutive items. The load is specialized for + //! ``BLOCK_LOAD_WARP_TRANSPOSE``, meaning memory references are efficiently coalesced using a warp-striped access + //! pattern (after which items are locally reordered among threads). + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, ...) + //! { + //! // Specialize BlockLoad for a 1D block of 128 threads owning 4 integer items each + //! using BlockLoad = cub::BlockLoad; + //! + //! // Allocate shared memory for BlockLoad + //! __shared__ typename BlockLoad::TempStorage temp_storage; + //! + //! // Load a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! BlockLoad(temp_storage).Load(d_data, thread_data); + //! } + //! + //! Suppose the input ``d_data`` is ``0, 1, 2, 3, 4, 5, ...``. The set of ``thread_data`` across the block of threads + //! in those threads will be ``{ [0,1,2,3], [4,5,6,7], ..., [508,509,510,511] }``. + //! + //! @endrst + //! + //! @param[in] block_src_it + //! The thread block's base iterator for loading from + //! + //! @param[out] dst_items + //! Destination to load data into + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Load(RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread]) + { + if constexpr (Algorithm == BLOCK_LOAD_DIRECT) + { + LoadDirectBlocked(linear_tid, block_src_it, dst_items); + } + else if constexpr (Algorithm == BLOCK_LOAD_STRIPED) + { + LoadDirectStriped(linear_tid, block_src_it, dst_items); + } + else if constexpr (Algorithm == BLOCK_LOAD_VECTORIZE) + { + if constexpr (detail::is_CacheModifiedInputIterator) + { + InternalLoadDirectBlockedVectorized(linear_tid, block_src_it.ptr, dst_items); + } + else if constexpr (::cuda::std::contiguous_iterator + && ::cuda::std::__can_to_address) + { + InternalLoadDirectBlockedVectorized(linear_tid, ::cuda::std::to_address(block_src_it), dst_items); + } + else + { + LoadDirectBlocked(linear_tid, block_src_it, dst_items); + } + } + else if constexpr (Algorithm == BLOCK_LOAD_TRANSPOSE) + { + LoadDirectStriped(linear_tid, block_src_it, dst_items); + block_exchange(temp_storage).StripedToBlocked(dst_items, dst_items); + } + else if constexpr (Algorithm == BLOCK_LOAD_WARP_TRANSPOSE || Algorithm == BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED) + { + LoadDirectWarpStriped(linear_tid, block_src_it, dst_items); + block_exchange(temp_storage).WarpStripedToBlocked(dst_items, dst_items); + } + } + + //! @rst + //! Load a linear segment of items from memory, guarded by range. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @blocked + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the guarded loading of a linear segment of 512 integers into a "blocked" + //! arrangement across 128 threads where each thread owns 4 consecutive items. The load is specialized for + //! ``BLOCK_LOAD_WARP_TRANSPOSE``, meaning memory references are efficiently coalesced using a warp-striped access + //! pattern (after which items are locally reordered among threads). + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, int block_items_end, ...) + //! { + //! // Specialize BlockLoad for a 1D block of 128 threads owning 4 integer items each + //! using BlockLoad = cub::BlockLoad; + //! + //! // Allocate shared memory for BlockLoad + //! __shared__ typename BlockLoad::TempStorage temp_storage; + //! + //! // Load a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! BlockLoad(temp_storage).Load(d_data, thread_data, block_items_end); + //! } + //! + //! Suppose the input ``d_data`` is ``0, 1, 2, 3, 4, 5, 6...`` and ``block_items_end`` is ``5``. The set of + //! ``thread_data`` across the block of threads in those threads will be ``{ [0,1,2,3], [4,?,?,?], ..., [?,?,?,?] }``, + //! with only the first two threads being unmasked to load portions of valid data (and other items remaining + //! unassigned). + //! + //! @endrst + //! + //! @param[in] block_src_it + //! The thread block's base iterator for loading from + //! + //! @param[out] dst_items + //! Destination to load data into + //! + //! @param[in] block_items_end + //! Number of valid items to load + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + Load(RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread], int block_items_end) + { + if constexpr (Algorithm == BLOCK_LOAD_DIRECT || Algorithm == BLOCK_LOAD_VECTORIZE) + { + LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end); + } + else if constexpr (Algorithm == BLOCK_LOAD_STRIPED) + { + LoadDirectStriped(linear_tid, block_src_it, dst_items, block_items_end); + } + else if constexpr (Algorithm == BLOCK_LOAD_TRANSPOSE) + { + LoadDirectStriped(linear_tid, block_src_it, dst_items, block_items_end); + block_exchange(temp_storage).StripedToBlocked(dst_items, dst_items); + } + else if constexpr (Algorithm == BLOCK_LOAD_WARP_TRANSPOSE || Algorithm == BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED) + { + LoadDirectWarpStriped(linear_tid, block_src_it, dst_items, block_items_end); + block_exchange(temp_storage).WarpStripedToBlocked(dst_items, dst_items); + } + } + + //! @rst + //! Load a linear segment of items from memory, guarded by range, with a fall-back assignment of out-of-bound elements + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @blocked + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the guarded loading of a linear segment of 512 integers into a "blocked" + //! arrangement across 128 threads where each thread owns 4 consecutive items. The load is specialized for + //! ``BLOCK_LOAD_WARP_TRANSPOSE``, meaning memory references are efficiently coalesced using a warp-striped access + //! pattern (after which items are locally reordered among threads). + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, int block_items_end, ...) + //! { + //! // Specialize BlockLoad for a 1D block of 128 threads owning 4 integer items each + //! using BlockLoad = cub::BlockLoad; + //! + //! // Allocate shared memory for BlockLoad + //! __shared__ typename BlockLoad::TempStorage temp_storage; + //! + //! // Load a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! BlockLoad(temp_storage).Load(d_data, thread_data, block_items_end, -1); + //! } + //! + //! Suppose the input ``d_data`` is ``0, 1, 2, 3, 4, 5, 6...``, ``block_items_end`` is ``5``, and the out-of-bounds + //! default is ``-1``. The set of ``thread_data`` across the block of threads in those threads will be + //! ``{ [0,1,2,3], [4,-1,-1,-1], ..., [-1,-1,-1,-1] }``, with only the first two threads being unmasked to load + //! portions of valid data (and other items are assigned ``-1``) + //! + //! @endrst + //! + //! @param[in] block_src_it + //! The thread block's base iterator for loading from + //! + //! @param[out] dst_items + //! Destination to load data into + //! + //! @param[in] block_items_end + //! Number of valid items to load + //! + //! @param[in] oob_default + //! Default value to assign out-of-bound items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + Load(RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread], int block_items_end, DefaultT oob_default) + { + if constexpr (Algorithm == BLOCK_LOAD_DIRECT || Algorithm == BLOCK_LOAD_VECTORIZE) + { + LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end, oob_default); + } + else if constexpr (Algorithm == BLOCK_LOAD_STRIPED) + { + LoadDirectStriped(linear_tid, block_src_it, dst_items, block_items_end, oob_default); + } + else if constexpr (Algorithm == BLOCK_LOAD_TRANSPOSE) + { + LoadDirectStriped(linear_tid, block_src_it, dst_items, block_items_end, oob_default); + block_exchange(temp_storage).StripedToBlocked(dst_items, dst_items); + } + else if constexpr (Algorithm == BLOCK_LOAD_WARP_TRANSPOSE || Algorithm == BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED) + { + LoadDirectWarpStriped(linear_tid, block_src_it, dst_items, block_items_end, oob_default); + block_exchange(temp_storage).WarpStripedToBlocked(dst_items, dst_items); + } + } + + //! @} +}; + +template > +struct BlockLoadType +{ + using type = cub::BlockLoad; +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_load_to_shared.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_load_to_shared.cuh new file mode 100644 index 00000000..cfbf64ba --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_load_to_shared.cuh @@ -0,0 +1,437 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! @file +//! The @c cub::BlockLoadToShared class provides a :ref:`collective ` method for asynchronously +//! loading data from global to shared memory. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +//! @rst +//! The @c BlockLoadToShared class provides a :ref:`collective ` method for asynchronously +//! loading data from global to shared memory. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - Given one or more spans of input elements in global memory and buffers in shared memory, this primitive +//! asynchronously copies the elements to shared memory and takes care of synchronization. +//! - @rowmajor +//! - Shared memory buffers are assumed to be sized according to `cub::detail::LoadToSharedBufferSize(num_items)` and aligned according to `cub::detail::LoadToSharedBufferAlignBytes()`. +//! - Global memory spans are by default assumed to be aligned according to the value type. Higher alignment guarantees +//! can optionally be specified. +//! - After one or more calls to `CopyAsync`, `Commit` needs to be called before optionally doing other work and then +//! calling `Wait` which guarantees the data to be available in shared memory, resets the state and allows for the +//! next call to `CopyAsync`. +//! +//! Performance Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - Uses special instructions/hardware acceleration when available (cp.async.bulk on Hopper+, copy.async on Ampere). +//! - By guaranteeing 16 byte alignment and size multiple for the global span, a faster path is taken and less shared +//! memory is needed for the destination buffer. +//! @endrst +template +struct BlockLoadToShared +{ +private: + /// Constants + static constexpr int threads_per_block = BlockDimX * BlockDimY * BlockDimZ; + + // Helper for fallback to gmem->reg->smem + struct alignas(detail::bulk_copy_min_align) vec_load_t + { + char c_array[detail::bulk_copy_min_align]; + }; + + struct _TempStorage + { + ::cuda::std::uint64_t mbarrier_handle; + }; + +#ifdef CCCL_ENABLE_DEVICE_ASSERTIONS + enum struct State + { + ready_to_copy, + ready_to_copy_or_commit, + committed, + invalidated, + }; +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS + + /// Shared storage reference + _TempStorage& temp_storage; + + const int linear_tid{cub::RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)}; + + // Thread selection for uniform operations + const bool elected{__elect_thread()}; + // Keep track of current mbarrier phase for waiting. + uint32_t phase_parity{}; + // Keep track of the amount of bytes from multiple transactions for Commit() (only needed for TMA). + // Also used to check for proper ordering of member function calls in debug mode. + uint32_t num_bytes_bulk_total{}; + +#ifdef CCCL_ENABLE_DEVICE_ASSERTIONS + State state{State::ready_to_copy}; +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS + + _CCCL_DEVICE_API _CCCL_FORCEINLINE bool __elect_thread() const + { + // Otherwise elect.sync in the last warp with a full mask is UB. + static_assert(threads_per_block % cub::detail::warp_threads == 0, + "The block size must be a multiple of the warp size"); + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_90, + ( // Use last warp to try to avoid having the elected thread also working on the peeling in the first warp. + return (linear_tid >= threads_per_block - cub::detail::warp_threads) && ::cuda::ptx::elect_sync(~0u);), + (return linear_tid == 0;)); + } + + _CCCL_DEVICE_API _CCCL_FORCEINLINE void __init_mbarrier() + { + { + NV_IF_TARGET(NV_PROVIDES_SM_90, ({ + if (elected) + { + ::cuda::ptx::mbarrier_init(&temp_storage.mbarrier_handle, 1); + } + // TODO The following sync was added to avoid a racecheck posititive. Is it really needed? + __syncthreads(); + })); + } + } + + _CCCL_DEVICE_API _CCCL_FORCEINLINE void __copy_aligned_async_bulk(char* smem_dst, const char* gmem_src, int num_bytes) + { + if (elected) + { +#if __cccl_ptx_isa >= 860 + NV_IF_TARGET(NV_PROVIDES_SM_90, ({ + ::cuda::ptx::cp_async_bulk( + ::cuda::ptx::space_shared, + ::cuda::ptx::space_global, + smem_dst, + gmem_src, + num_bytes, + &temp_storage.mbarrier_handle); + })); +#else + NV_IF_TARGET(NV_PROVIDES_SM_90, ({ + ::cuda::ptx::cp_async_bulk( + ::cuda::ptx::space_cluster, + ::cuda::ptx::space_global, + smem_dst, + gmem_src, + num_bytes, + &temp_storage.mbarrier_handle); + })); +#endif // __cccl_ptx_isa >= 800 + // Needed for arrival on mbarrier in Commit() + num_bytes_bulk_total += num_bytes; + } + } + + _CCCL_DEVICE_API _CCCL_FORCEINLINE void __copy_aligned_async(char* smem_dst, const char* gmem_src, int num_bytes) + { + for (int offset = linear_tid * detail::bulk_copy_min_align; offset < num_bytes; + offset += threads_per_block * detail::bulk_copy_min_align) + { + [[maybe_unused]] const auto thread_src = gmem_src + offset; + [[maybe_unused]] const auto thread_dst = smem_dst + offset; + // LDGSTS borrowed from cuda::memcpy_async, assumes 16 byte alignment to avoid L1 (.cg) + NV_IF_TARGET( + NV_PROVIDES_SM_80, ({ + asm volatile( + "cp.async.cg.shared.global [%0], [%1], %2, %2;" + : + : "r"(static_cast<::cuda::std::uint32_t>(::__cvta_generic_to_shared(thread_dst))), "l"(thread_src), "n"(16) + : "memory"); + })); + } + } + + _CCCL_DEVICE_API _CCCL_FORCEINLINE void __copy_aligned_fallback(char* smem_dst, const char* gmem_src, int num_bytes) + { + for (int offset = linear_tid * detail::bulk_copy_min_align; offset < num_bytes; + offset += threads_per_block * detail::bulk_copy_min_align) + { + const auto thread_src = gmem_src + offset; + const auto thread_dst = smem_dst + offset; + *::cuda::ptr_rebind(thread_dst) = *::cuda::ptr_rebind(thread_src); + } + } + + _CCCL_DEVICE_API _CCCL_FORCEINLINE void __copy_aligned(char* smem_dst, const char* gmem_src, int num_bytes) + { + NV_DISPATCH_TARGET( + NV_PROVIDES_SM_90, + (__copy_aligned_async_bulk(smem_dst, gmem_src, num_bytes);), + NV_PROVIDES_SM_80, + (__copy_aligned_async(smem_dst, gmem_src, num_bytes);), + NV_IS_DEVICE, + (__copy_aligned_fallback(smem_dst, gmem_src, num_bytes);)); + } + + // Dispatch to fallback for waiting pre TMA/SM_90 + _CCCL_DEVICE_API _CCCL_FORCEINLINE bool __try_wait() + { + NV_DISPATCH_TARGET( + NV_PROVIDES_SM_90, + (return ::cuda::ptx::mbarrier_try_wait_parity(&temp_storage.mbarrier_handle, phase_parity);), + NV_PROVIDES_SM_80, + (asm volatile("cp.async.wait_group 0;" :: : "memory"); // + __syncthreads(); + return true;), + NV_ANY_TARGET, + (__syncthreads(); // + return true;)); + } + + // token is only constructible by BlockLoadToShared + class token_impl + { + friend struct BlockLoadToShared; + _CCCL_DEVICE_API _CCCL_FORCEINLINE token_impl() {} // NOLINT(modernize-use-equals-default) ctor must have a body to + // avoid token_impl{} to compile + + public: + // NOLINTBEGIN(modernize-use-equals-delete) + token_impl(const token_impl&) = delete; + token_impl& operator=(const token_impl&) = delete; + // NOLINTEND(modernize-use-equals-delete) + + token_impl(token_impl&&) = default; + token_impl& operator=(token_impl&&) = default; + }; + +public: + /// @smemstorage{BlockLoadToShared} + using TempStorage = cub::Uninitialized<_TempStorage>; + + //! Token type used to enforce correct call order between Commit() and Wait() + //! member functions. Returned by Commit() and required by Wait() as a usage + //! guard. + using CommitToken = token_impl; + + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using the specified memory allocation as temporary storage. + //! + //! @param[in] temp_storage + //! Reference to memory allocation having layout type TempStorage + _CCCL_DEVICE_API _CCCL_FORCEINLINE BlockLoadToShared(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + { + _CCCL_ASSERT(::cuda::device::is_object_from(temp_storage, ::cuda::device::address_space::shared), + "temp_storage has to be in shared memory"); + __init_mbarrier(); + } + + _CCCL_DEVICE_API BlockLoadToShared(const BlockLoadToShared&) = delete; + + //! @} + + _CCCL_DEVICE_API BlockLoadToShared& operator=(const BlockLoadToShared&) = delete; + + //! @brief Invalidates underlying @c mbarrier enabling reuse of its temporary storage. + //! @note + //! Block-synchronization is needed after calling `Invalidate()` to reuse the shared memory from the temporary + //! storage. + // This is not the destructor to avoid overhead when shared memory reuse is not needed. + _CCCL_DEVICE_API _CCCL_FORCEINLINE void Invalidate() + { +#ifdef CCCL_ENABLE_DEVICE_ASSERTIONS + _CCCL_ASSERT(state == State::ready_to_copy, "Wait() must be called before Invalidate()"); + state = State::invalidated; +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS + // Make sure all threads are done interacting with the mbarrier + __syncthreads(); + if (elected) + { + NV_IF_TARGET(NV_PROVIDES_SM_90, ::cuda::ptx::mbarrier_inval(&temp_storage.mbarrier_handle);); + } + // Make sure the elected thread is done invalidating the mbarrier + __syncthreads(); + } + + //! @brief Copy elements from global to shared memory + //! @tparam T + //! **[inferred]** Value type for this transaction + //! @tparam GmemAlign + //! Guaranteed alignment in bytes of the source range (both begin and end) in global memory + //! @param[in] smem_dst + //! Destination buffer in shared memory that is aligned to `SharedBufferAlignBytes()` and at least + //! `SharedBufferSizeBytes(size(gmem_src))` big. + //! @param[in] gmem_src + //! Source range in global memory, determines the size of the transaction + //! @return + //! The range in shared memory (same size as `gmem_src`) which should be used to access the data after `Commit` and + //! `Wait`. + //! Note: This range is aliasing the `smem_dst` buffer. So `smem_dst` should not be written to/reused while this + //! range is still in use! + // TODO Allow spans with static sizes? + template + [[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE ::cuda::std::span + CopyAsync(::cuda::std::span smem_dst, ::cuda::std::span gmem_src) + { + static_assert(THRUST_NS_QUALIFIER::is_trivially_relocatable_v); + static_assert(::cuda::__is_valid_alignment(GmemAlign)); + constexpr bool bulk_aligned = GmemAlign >= static_cast<::cuda::std::size_t>(detail::bulk_copy_min_align); + // Avoid 64b multiplication in span::size_bytes() + const int num_bytes = static_cast(sizeof(T)) * static_cast(size(gmem_src)); + const auto dst_ptr = data(smem_dst); + const auto src_ptr = ::cuda::ptr_rebind(data(gmem_src)); + _CCCL_ASSERT(dst_ptr == nullptr || ::cuda::device::is_address_from(dst_ptr, ::cuda::device::address_space::shared), + "Destination address needs to point to shared memory"); + _CCCL_ASSERT(src_ptr == nullptr || ::cuda::device::is_address_from(src_ptr, ::cuda::device::address_space::global), + "Source address needs to point to global memory"); + _CCCL_ASSERT((src_ptr != nullptr && dst_ptr != nullptr) || num_bytes == 0, + "Only when the source range is empty are nullptrs allowed"); + _CCCL_ASSERT(::cuda::is_aligned(src_ptr, GmemAlign), + "Begin of global memory range needs to be aligned according to GmemAlign."); + _CCCL_ASSERT(::cuda::is_aligned(src_ptr + num_bytes, GmemAlign), + "End of global memory range needs to be aligned according to GmemAlign."); + _CCCL_ASSERT(::cuda::is_aligned(dst_ptr, cub::detail::LoadToSharedBufferAlignBytes()), + "Shared memory needs to be 16 byte aligned."); + _CCCL_ASSERT( + (static_cast(size(smem_dst)) >= cub::detail::LoadToSharedBufferSizeBytes(size(gmem_src))), + "Shared memory destination buffer must have enough space"); +#ifdef CCCL_ENABLE_DEVICE_ASSERTIONS + _CCCL_ASSERT(state == State::ready_to_copy || state == State::ready_to_copy_or_commit, + "Wait() must be called before another CopyAsync()"); + state = State::ready_to_copy_or_commit; +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS + if constexpr (bulk_aligned) + { + __copy_aligned(dst_ptr, src_ptr, num_bytes); + return {::cuda::ptr_rebind(::cuda::std::data(smem_dst)), ::cuda::std::size(gmem_src)}; + } + else + { + const auto src_ptr_aligned = ::cuda::align_up(src_ptr, detail::bulk_copy_min_align); + const int align_diff = static_cast(src_ptr_aligned - src_ptr); + const int head_padding_bytes = (detail::bulk_copy_min_align - align_diff) % detail::bulk_copy_min_align; + const auto actual_dst_ptr = dst_ptr + head_padding_bytes; + const int head_peeling_bytes = ::cuda::std::min(align_diff, num_bytes); + const int num_bytes_bulk = ::cuda::round_down(num_bytes - head_peeling_bytes, detail::bulk_copy_min_align); + __copy_aligned(actual_dst_ptr + head_peeling_bytes, src_ptr_aligned, num_bytes_bulk); + + // Peel head and tail + // Make sure we have enough threads for the worst case of bulk_min_align bytes on each side. + static_assert(threads_per_block >= 2 * (detail::bulk_copy_min_align - 1)); + // |-------------head--------------|--------------------------tail--------------------------| + // 0, 1, ... head_peeling_bytes - 1, head_peeling_bytes + num_bytes_bulk, ..., num_bytes - 1 + const int begin_offset = linear_tid < head_peeling_bytes ? 0 : num_bytes_bulk; + if (const int idx = begin_offset + linear_tid; idx < num_bytes) + { + actual_dst_ptr[idx] = src_ptr[idx]; + } + return {::cuda::ptr_rebind(actual_dst_ptr), ::cuda::std::size(gmem_src)}; + } + } + + // Avoid need to explicitly specify `T` for non-const src. + //! @brief Convenience overload, see `CopyAsync(span, span)`. + template + [[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE ::cuda::std::span + CopyAsync(::cuda::std::span smem_dst, ::cuda::std::span gmem_src) + { + return CopyAsync(smem_dst, ::cuda::std::span{gmem_src}); + } + + //! @brief Commit one or more @c CopyAsync() calls. + [[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE CommitToken Commit() + { +#ifdef CCCL_ENABLE_DEVICE_ASSERTIONS + _CCCL_ASSERT(state == State::ready_to_copy_or_commit, "CopyAsync() must be called before Commit()"); + state = State::committed; +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS + + NV_DISPATCH_TARGET( + NV_PROVIDES_SM_90, + (if (elected) { + ::cuda::ptx::mbarrier_arrive_expect_tx( + ::cuda::ptx::sem_release, + ::cuda::ptx::scope_cta, + ::cuda::ptx::space_shared, + &temp_storage.mbarrier_handle, + num_bytes_bulk_total); + num_bytes_bulk_total = 0u; + } // + __syncthreads();), + NV_PROVIDES_SM_80, + (asm volatile("cp.async.commit_group ;" :: : "memory");)); + + // Token's mere purpose currently is to prevent calling Wait() without a + // prior Commit() + return CommitToken{}; + } + + //! @brief Wait for previously committed copies to arrive. Prepare for next + //! calls to @c CopyAsync() . + _CCCL_DEVICE_API _CCCL_FORCEINLINE void Wait(CommitToken&&) + { +#ifdef CCCL_ENABLE_DEVICE_ASSERTIONS + _CCCL_ASSERT(state == State::committed, "Commit() must be called before Wait()"); + state = State::ready_to_copy; +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS + + while (!__try_wait()) + ; + phase_parity ^= 1u; + } + + //! @brief Convenience overload calling `Commit()` and `Wait()`. + _CCCL_DEVICE_API _CCCL_FORCEINLINE void CommitAndWait() + { + Wait(Commit()); + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_merge_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_merge_sort.cuh new file mode 100644 index 00000000..b1b00e3a --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_merge_sort.cuh @@ -0,0 +1,859 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! Computes the intersection of the diagonal \c diag with the merge path in the merge matrix of two input sequences. +//! This implements the DiagonalIntersection algorithm from Merge-Path. Additional details can be found in: +//! * S. Odeh, O. Green, Z. Mwassi, O. Shmueli, Y. Birk, "Merge Path - Parallel Merging Made Simple", Multithreaded +//! Architectures and Applications (MTAAP) Workshop, IEEE 26th International Parallel & Distributed Processing +//! Symposium (IPDPS), 2012 +//! * S. Odeh, O. Green, Y. Birk, "Merge Path - A Visually Intuitive Approach to Parallel Merging", 2014, URL: +//! https://arxiv.org/abs/1406.2628 +//! \returns The number of elements merged from the first sequence at the intersection of the diagonal with the merge +//! path. The number of elements merged from the second sequence is \c diag minus the returned value. +template +_CCCL_DEVICE _CCCL_FORCEINLINE OffsetT +MergePath(KeyIt1 keys1, KeyIt2 keys2, OffsetT keys1_count, OffsetT keys2_count, OffsetT diag, BinaryPred binary_pred) +{ + OffsetT keys1_begin = diag < keys2_count ? 0 : diag - keys2_count; + OffsetT keys1_end = (::cuda::std::min) (diag, keys1_count); + + while (keys1_begin < keys1_end) + { + const OffsetT mid = cub::MidPoint(keys1_begin, keys1_end); + // pull copies of the keys before calling binary_pred so proxy references are unwrapped + const detail::it_value_t key1 = keys1[mid]; + const detail::it_value_t key2 = keys2[diag - 1 - mid]; + if (binary_pred(key2, key1)) + { + keys1_end = mid; + } + else + { + keys1_begin = mid + 1; + } + } + return keys1_begin; +} + +namespace detail +{ +template +_CCCL_DEVICE _CCCL_FORCEINLINE void serial_merge( + KeyIt keys_shared, + int keys1_beg, + int keys2_beg, + int keys1_count, + int keys2_count, + KeyT (&output)[ItemsPerThread], + int (&indices)[ItemsPerThread], + CompareOp compare_op, + KeyT oob_default) +{ + const int keys1_end = keys1_beg + keys1_count; + const int keys2_end = keys2_beg + keys2_count; + + KeyT key1 = keys1_count != 0 ? keys_shared[keys1_beg] : oob_default; + KeyT key2 = keys2_count != 0 ? keys_shared[keys2_beg] : oob_default; + + _CCCL_PRAGMA_UNROLL(Unroll ? ItemsPerThread : 1) + for (int item = 0; item < ItemsPerThread; ++item) + { + const bool p = (keys2_beg < keys2_end) && ((keys1_beg >= keys1_end) || compare_op(key2, key1)); + output[item] = p ? key2 : key1; + indices[item] = p ? keys2_beg++ : keys1_beg++; + if (p) + { + key2 = keys_shared[keys2_beg]; + } + else + { + key1 = keys_shared[keys1_beg]; + } + } +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void serial_merge( + KeyIt keys_shared, + int keys1_beg, + int keys2_beg, + int keys1_count, + int keys2_count, + KeyT (&output)[ItemsPerThread], + int (&indices)[ItemsPerThread], + CompareOp compare_op) +{ + serial_merge( + keys_shared, keys1_beg, keys2_beg, keys1_count, keys2_count, output, indices, compare_op, output[0]); +} +} // namespace detail + +//! Merges elements from two sorted sequences +//! \tparam ItemsPerThread The number of elements to merge and write to \c output +//! \param keys_shared An iterator to shared memory containing from which both sequences are reachable +//! \param keys1_beg The index into \c keys_shared where the first sequence starts +//! \param keys2_beg The index into \c keys_shared where the second sequence starts +//! \param keys1_count The maximum number of keys to merge from the first sequence. One more item may be read but is not +//! used. +//! \param keys2_count The maximum number of keys to merge from the second sequence. One more item may be read but is +//! not used. +//! \param output The output array +//! \param indices The shared memory indices relative to \c keys_shared of the elements written to \c output +template +_CCCL_DEVICE _CCCL_FORCEINLINE void SerialMerge( + KeyIt keys_shared, + int keys1_beg, + int keys2_beg, + int keys1_count, + int keys2_count, + KeyT (&output)[ItemsPerThread], + int (&indices)[ItemsPerThread], + CompareOp compare_op, + KeyT oob_default) +{ + detail::serial_merge( + keys_shared, keys1_beg, keys2_beg, keys1_count, keys2_count, output, indices, compare_op, oob_default); +} + +template +_CCCL_DEVICE _CCCL_FORCEINLINE void SerialMerge( + KeyIt keys_shared, + int keys1_beg, + int keys2_beg, + int keys1_count, + int keys2_count, + KeyT (&output)[ItemsPerThread], + int (&indices)[ItemsPerThread], + CompareOp compare_op) +{ + detail::serial_merge(keys_shared, keys1_beg, keys2_beg, keys1_count, keys2_count, output, indices, compare_op); +} + +/** + * @brief Generalized merge sort algorithm + * + * This class is used to reduce code duplication. Warp and Block merge sort + * differ only in how they compute thread index and how they synchronize + * threads. Since synchronization might require access to custom data + * (like member mask), CRTP is used. + * + * @par + * The code snippet below illustrates the way this class can be used. + * @par + * @code + * #include // or equivalently + * + * constexpr int BLOCK_THREADS = 256; + * constexpr int ItemsPerThread = 9; + * + * class BlockMergeSort : public BlockMergeSortStrategy + * { + * using BlockMergeSortStrategyT = + * BlockMergeSortStrategy; + * public: + * __device__ __forceinline__ explicit BlockMergeSort( + * typename BlockMergeSortStrategyT::TempStorage &temp_storage) + * : BlockMergeSortStrategyT(temp_storage, threadIdx.x) + * {} + * + * __device__ __forceinline__ void SyncImplementation() const + * { + * __syncthreads(); + * } + * }; + * @endcode + * + * @tparam KeyT + * KeyT type + * + * @tparam ValueT + * ValueT type. cub::NullType indicates a keys-only sort + * + * @tparam SynchronizationPolicy + * Provides a way of synchronizing threads. Should be derived from + * `BlockMergeSortStrategy`. + */ +template +class BlockMergeSortStrategy +{ + static_assert(::cuda::is_power_of_two(NumThreads), "NumThreads must be a power of two"); + +private: + static constexpr int ITEMS_PER_TILE = ItemsPerThread * NumThreads; + + // Whether or not there are values to be trucked along with keys + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + /// Shared memory type required by this thread block + union _TempStorage + { + KeyT keys_shared[ITEMS_PER_TILE + 1]; + ValueT items_shared[ITEMS_PER_TILE + 1]; + }; // union TempStorage +#endif // _CCCL_DOXYGEN_INVOKED + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + const unsigned int linear_tid; + +public: + /// \smemstorage{BlockMergeSort} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + BlockMergeSortStrategy() = delete; + explicit _CCCL_DEVICE _CCCL_FORCEINLINE BlockMergeSortStrategy(unsigned int linear_tid) + : temp_storage(PrivateStorage()) + , linear_tid(linear_tid) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE BlockMergeSortStrategy(TempStorage& temp_storage, unsigned int linear_tid) + : temp_storage(temp_storage.Alias()) + , linear_tid(linear_tid) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int get_linear_tid() const + { + return linear_tid; + } + + /** + * @brief Sorts items partitioned across a CUDA thread block using + * a merge sorting method. + * + * @par + * Sort is not guaranteed to be stable. That is, suppose that i and j are + * equivalent: neither one is less than the other. It is not guaranteed + * that the relative order of these two elements will be preserved by sort. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @tparam CompareOp + * functor type having member `bool operator()(KeyT lhs, KeyT rhs)`. + * `CompareOp` is a model of [Strict Weak Ordering]. + * + * @param[in,out] keys + * Keys to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Sort(KeyT (&keys)[ItemsPerThread], CompareOp compare_op) + { + ValueT items[ItemsPerThread]; + Sort(keys, items, compare_op, ITEMS_PER_TILE, keys[0]); + } + + /** + * @brief Sorts items partitioned across a CUDA thread block using + * a merge sorting method. + * + * @par + * - Sort is not guaranteed to be stable. That is, suppose that `i` and `j` + * are equivalent: neither one is less than the other. It is not guaranteed + * that the relative order of these two elements will be preserved by sort. + * - The value of `oob_default` is assigned to all elements that are out of + * `valid_items` boundaries. It's expected that `oob_default` is ordered + * after any value in the `valid_items` boundaries. The algorithm always + * sorts a fixed amount of elements, which is equal to + * `ItemsPerThread * BLOCK_THREADS`. If there is a value that is ordered + * after `oob_default`, it won't be placed within `valid_items` boundaries. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @tparam CompareOp + * functor type having member `bool operator()(KeyT lhs, KeyT rhs)`. + * `CompareOp` is a model of [Strict Weak Ordering]. + * + * @param[in,out] keys + * Keys to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] valid_items + * Number of valid items to sort + * + * @param[in] oob_default + * Default value to assign out-of-bound items + * + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + Sort(KeyT (&keys)[ItemsPerThread], CompareOp compare_op, int valid_items, KeyT oob_default) + { + ValueT items[ItemsPerThread]; + Sort(keys, items, compare_op, valid_items, oob_default); + } + + /** + * @brief Sorts items partitioned across a CUDA thread block using a merge sorting method. + * + * @par + * Sort is not guaranteed to be stable. That is, suppose that `i` and `j` are + * equivalent: neither one is less than the other. It is not guaranteed + * that the relative order of these two elements will be preserved by sort. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @tparam CompareOp + * functor type having member `bool operator()(KeyT lhs, KeyT rhs)`. + * `CompareOp` is a model of [Strict Weak Ordering]. + * + * @param[in,out] keys + * Keys to sort + * + * @param[in,out] items + * Values to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + Sort(KeyT (&keys)[ItemsPerThread], ValueT (&items)[ItemsPerThread], CompareOp compare_op) + { + Sort(keys, items, compare_op, ITEMS_PER_TILE, keys[0]); + } + + /** + * @brief Sorts items partitioned across a CUDA thread block using + * a merge sorting method. + * + * @par + * - Sort is not guaranteed to be stable. That is, suppose that `i` and `j` + * are equivalent: neither one is less than the other. It is not guaranteed + * that the relative order of these two elements will be preserved by sort. + * - The value of `oob_default` is assigned to all elements that are out of + * `valid_items` boundaries. It's expected that `oob_default` is ordered + * after any value in the `valid_items` boundaries. The algorithm always + * sorts a fixed amount of elements, which is equal to + * `ItemsPerThread * BLOCK_THREADS`. If there is a value that is ordered + * after `oob_default`, it won't be placed within `valid_items` boundaries. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @tparam CompareOp + * functor type having member `bool operator()(KeyT lhs, KeyT rhs)` + * `CompareOp` is a model of [Strict Weak Ordering]. + * + * @tparam IS_LAST_TILE + * True if `valid_items` isn't equal to the `ITEMS_PER_TILE` + * + * @param[in,out] keys + * Keys to sort + * + * @param[in,out] items + * Values to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] valid_items + * Number of valid items to sort + * + * @param[in] oob_default + * Default value to assign out-of-bound items + * + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + Sort(KeyT (&keys)[ItemsPerThread], + ValueT (&items)[ItemsPerThread], + CompareOp compare_op, + int valid_items, + KeyT oob_default) + { + if constexpr (IS_LAST_TILE) + { + // if last tile, find valid max_key + // and fill the remaining keys with it + // + KeyT max_key = oob_default; + + _CCCL_PRAGMA_UNROLL(_Unroll ? ItemsPerThread : 1) + for (int item = 1; item < ItemsPerThread; ++item) + { + if (ItemsPerThread * linear_tid + item < valid_items) + { + max_key = compare_op(max_key, keys[item]) ? keys[item] : max_key; + } + else + { + keys[item] = max_key; + } + } + } + + // if first element of thread is in input range, stable sort items + // + if (!IS_LAST_TILE || ItemsPerThread * linear_tid < valid_items) + { + detail::stable_odd_even_sort<_Unroll>(keys, items, compare_op); + } + + // each thread has sorted keys + // merge sort keys in shared memory + // + for (int target_merged_threads_number = 2; target_merged_threads_number <= NumThreads; + target_merged_threads_number *= 2) + { + const int merged_threads_number = target_merged_threads_number / 2; + const int mask = target_merged_threads_number - 1; + + Sync(); + + // store keys in shmem + // + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ItemsPerThread; ++item) + { + int idx = ItemsPerThread * linear_tid + item; + temp_storage.keys_shared[idx] = keys[item]; + } + + Sync(); + + int indices[ItemsPerThread]; + + const int first_thread_idx_in_thread_group_being_merged = ~mask & linear_tid; + const int start = ItemsPerThread * first_thread_idx_in_thread_group_being_merged; + const int size = ItemsPerThread * merged_threads_number; + + const int thread_idx_in_thread_group_being_merged = mask & linear_tid; + + const int diag = (::cuda::std::min) (valid_items, ItemsPerThread * thread_idx_in_thread_group_being_merged); + + const int keys1_beg = (::cuda::std::min) (valid_items, start); + const int keys1_end = (::cuda::std::min) (valid_items, keys1_beg + size); + const int keys2_beg = keys1_end; + const int keys2_end = (::cuda::std::min) (valid_items, keys2_beg + size); + + const int keys1_count = keys1_end - keys1_beg; + const int keys2_count = keys2_end - keys2_beg; + + const int partition_diag = MergePath( + &temp_storage.keys_shared[keys1_beg], + &temp_storage.keys_shared[keys2_beg], + keys1_count, + keys2_count, + diag, + compare_op); + + const int keys1_beg_loc = keys1_beg + partition_diag; + const int keys1_end_loc = keys1_end; + const int keys2_beg_loc = keys2_beg + diag - partition_diag; + const int keys2_end_loc = keys2_end; + const int keys1_count_loc = keys1_end_loc - keys1_beg_loc; + const int keys2_count_loc = keys2_end_loc - keys2_beg_loc; + detail::serial_merge<_Unroll>( + &temp_storage.keys_shared[0], + keys1_beg_loc, + keys2_beg_loc, + keys1_count_loc, + keys2_count_loc, + keys, + indices, + compare_op, + oob_default); + + if constexpr (!KEYS_ONLY) + { + Sync(); + + // store keys in shmem + // + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ItemsPerThread; ++item) + { + int idx = ItemsPerThread * linear_tid + item; + temp_storage.items_shared[idx] = items[item]; + } + + Sync(); + + // gather items from shmem + // + _CCCL_PRAGMA_UNROLL_FULL() + for (int item = 0; item < ItemsPerThread; ++item) + { + items[item] = temp_storage.items_shared[indices[item]]; + } + } + } + } // func block_merge_sort + + /** + * @brief Sorts items partitioned across a CUDA thread block using + * a merge sorting method. + * + * @par + * StableSort is stable: it preserves the relative ordering of equivalent + * elements. That is, if `x` and `y` are elements such that `x` precedes `y`, + * and if the two elements are equivalent (neither `x < y` nor `y < x`) then + * a postcondition of StableSort is that `x` still precedes `y`. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @tparam CompareOp + * functor type having member `bool operator()(KeyT lhs, KeyT rhs)`. + * `CompareOp` is a model of [Strict Weak Ordering]. + * + * @param[in,out] keys + * Keys to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void StableSort(KeyT (&keys)[ItemsPerThread], CompareOp compare_op) + { + Sort(keys, compare_op); + } + + /** + * @brief Sorts items partitioned across a CUDA thread block using + * a merge sorting method. + * + * @par + * StableSort is stable: it preserves the relative ordering of equivalent + * elements. That is, if `x` and `y` are elements such that `x` precedes `y`, + * and if the two elements are equivalent (neither `x < y` nor `y < x`) then + * a postcondition of StableSort is that `x` still precedes `y`. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @tparam CompareOp + * functor type having member `bool operator()(KeyT lhs, KeyT rhs)`. + * `CompareOp` is a model of [Strict Weak Ordering]. + * + * @param[in,out] keys + * Keys to sort + * + * @param[in,out] items + * Values to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + StableSort(KeyT (&keys)[ItemsPerThread], ValueT (&items)[ItemsPerThread], CompareOp compare_op) + { + Sort(keys, items, compare_op); + } + + /** + * @brief Sorts items partitioned across a CUDA thread block using + * a merge sorting method. + * + * @par + * - StableSort is stable: it preserves the relative ordering of equivalent + * elements. That is, if `x` and `y` are elements such that `x` precedes + * `y`, and if the two elements are equivalent (neither `x < y` nor `y < x`) + * then a postcondition of StableSort is that `x` still precedes `y`. + * - The value of `oob_default` is assigned to all elements that are out of + * `valid_items` boundaries. It's expected that `oob_default` is ordered + * after any value in the `valid_items` boundaries. The algorithm always + * sorts a fixed amount of elements, which is equal to + * `ItemsPerThread * BLOCK_THREADS`. + * If there is a value that is ordered after `oob_default`, it won't be + * placed within `valid_items` boundaries. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @tparam CompareOp + * functor type having member `bool operator()(KeyT lhs, KeyT rhs)`. + * `CompareOp` is a model of [Strict Weak Ordering]. + * + * @param[in,out] keys + * Keys to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] valid_items + * Number of valid items to sort + * + * @param[in] oob_default + * Default value to assign out-of-bound items + * + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + StableSort(KeyT (&keys)[ItemsPerThread], CompareOp compare_op, int valid_items, KeyT oob_default) + { + Sort(keys, compare_op, valid_items, oob_default); + } + + /** + * @brief Sorts items partitioned across a CUDA thread block using + * a merge sorting method. + * + * @par + * - StableSort is stable: it preserves the relative ordering of equivalent + * elements. That is, if `x` and `y` are elements such that `x` precedes + * `y`, and if the two elements are equivalent (neither `x < y` nor `y < x`) + * then a postcondition of StableSort is that `x` still precedes `y`. + * - The value of `oob_default` is assigned to all elements that are out of + * `valid_items` boundaries. It's expected that `oob_default` is ordered + * after any value in the `valid_items` boundaries. The algorithm always + * sorts a fixed amount of elements, which is equal to + * `ItemsPerThread * BLOCK_THREADS`. If there is a value that is ordered + * after `oob_default`, it won't be placed within `valid_items` boundaries. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @tparam CompareOp + * functor type having member `bool operator()(KeyT lhs, KeyT rhs)`. + * `CompareOp` is a model of [Strict Weak Ordering]. + * + * @tparam IS_LAST_TILE + * True if `valid_items` isn't equal to the `ITEMS_PER_TILE` + * + * @param[in,out] keys + * Keys to sort + * + * @param[in,out] items + * Values to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] valid_items + * Number of valid items to sort + * + * @param[in] oob_default + * Default value to assign out-of-bound items + * + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void StableSort( + KeyT (&keys)[ItemsPerThread], + ValueT (&items)[ItemsPerThread], + CompareOp compare_op, + int valid_items, + KeyT oob_default) + { + Sort(keys, items, compare_op, valid_items, oob_default); + } + +private: + _CCCL_DEVICE _CCCL_FORCEINLINE void Sync() const + { + static_cast(this)->SyncImplementation(); + } +}; + +/** + * @brief The BlockMergeSort class provides methods for sorting items + * partitioned across a CUDA thread block using a merge sorting method. + * + * @tparam KeyT + * KeyT type + * + * @tparam BLOCK_DIM_X + * The thread block length in threads along the X dimension + * + * @tparam ItemsPerThread + * The number of items per thread + * + * @tparam ValueT + * **[optional]** ValueT type (default: `cub::NullType`, which indicates + * a keys-only sort) + * + * @tparam BLOCK_DIM_Y + * **[optional]** The thread block length in threads along the Y dimension + * (default: 1) + * + * @tparam BLOCK_DIM_Z + * **[optional]** The thread block length in threads along the Z dimension + * (default: 1) + * + * @par Overview + * BlockMergeSort arranges items into ascending order using a comparison + * functor with less-than semantics. Merge sort can handle arbitrary types + * and comparison functors, but is slower than BlockRadixSort when sorting + * arithmetic types into ascending/descending order. + * + * @par A Simple Example + * @blockcollective{BlockMergeSort} + * @par + * The code snippet below illustrates a sort of 512 integer keys that are + * partitioned across 128 threads * where each thread owns 4 consecutive items. + * @par + * @code + * #include // or equivalently + * + * struct CustomLess + * { + * template + * __device__ bool operator()(const DataType &lhs, const DataType &rhs) + * { + * return lhs < rhs; + * } + * }; + * + * __global__ void ExampleKernel(...) + * { + * // Specialize BlockMergeSort for a 1D block of 128 threads owning 4 integer items each + * using BlockMergeSort = cub::BlockMergeSort; + * + * // Allocate shared memory for BlockMergeSort + * __shared__ typename BlockMergeSort::TempStorage temp_storage_shuffle; + * + * // Obtain a segment of consecutive items that are blocked across threads + * int thread_keys[4]; + * ... + * + * BlockMergeSort(temp_storage_shuffle).Sort(thread_keys, CustomLess()); + * ... + * } + * @endcode + * @par + * Suppose the set of input `thread_keys` across the block of threads is + * `{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }`. + * The corresponding output `thread_keys` in those threads will be + * `{ [0,1,2,3], [4,5,6,7], [8,9,10,11], ..., [508,509,510,511] }`. + * + * @par Re-using dynamically allocating shared memory + * The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of + * dynamically shared memory with BlockReduce and how to re-purpose + * the same memory region. + * + * This example can be easily adapted to the storage required by BlockMergeSort. + */ +template +class BlockMergeSort + : public BlockMergeSortStrategy< + KeyT, + ValueT, + BlockDimX * BlockDimY * BlockDimZ, + ItemsPerThread, + BlockMergeSort, + _Unroll> +{ +private: + // The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + static constexpr int ITEMS_PER_TILE = ItemsPerThread * BLOCK_THREADS; + + using BlockMergeSortStrategyT = + BlockMergeSortStrategy; + +public: + _CCCL_DEVICE _CCCL_FORCEINLINE BlockMergeSort() + : BlockMergeSortStrategyT(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE explicit BlockMergeSort(typename BlockMergeSortStrategyT::TempStorage& temp_storage) + : BlockMergeSortStrategyT(temp_storage, RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + +private: + _CCCL_DEVICE _CCCL_FORCEINLINE void SyncImplementation() const + { + __syncthreads(); + } + + friend BlockMergeSortStrategyT; +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_radix_rank.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_radix_rank.cuh new file mode 100644 index 00000000..b8bffa38 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_radix_rank.cuh @@ -0,0 +1,1243 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::BlockRadixRank provides operations for ranking unsigned integer types within a CUDA thread block + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @brief Radix ranking algorithm, the algorithm used to implement stable ranking of the +//! keys from a single tile. Note that different ranking algorithms require different +//! initial arrangements of keys to function properly. +enum RadixRankAlgorithm +{ + //! Ranking using the BlockRadixRank algorithm with `MemoizeOuterScan == false`. + //! It uses thread-private histograms, and thus uses more shared memory. + //! Requires blocked arrangement of keys. Does not support count callbacks. + RADIX_RANK_BASIC, + + //! Ranking using the BlockRadixRank algorithm with `MemoizeOuterScan == true`. + //! Similar to RADIX_RANK BASIC, it requires blocked arrangement of keys and does not support count callbacks. + RADIX_RANK_MEMOIZE, + + //! Ranking using the BlockRadixRankMatch algorithm. It uses warp-private histograms and matching for ranking + //! the keys in a single warp. Therefore, it uses less shared memory compared to RADIX_RANK_BASIC. + //! It requires warp-striped key arrangement and supports count callbacks. + RADIX_RANK_MATCH, + + //! Ranking using the BlockRadixRankMatchEarlyCounts algorithm with `MATCH_ALGORITHM == WARP_MATCH_ANY`. + //! An alternative implementation of match-based ranking that computes bin counts early. + //! Because of this, it works better with onesweep sorting, which requires bin counts for decoupled look-back. + //! Assumes warp-striped key arrangement and supports count callbacks. + RADIX_RANK_MATCH_EARLY_COUNTS_ANY, + + //! Ranking using the BlockRadixRankEarlyCounts algorithm with `MATCH_ALGORITHM == WARP_MATCH_ATOMIC_OR`. + //! It uses extra space in shared memory to generate warp match masks using `atomicOr()`. + //! This is faster when there are few matches, but can lead to slowdowns if the number of matching keys among + //! warp lanes is high. Assumes warp-striped key arrangement and supports count callbacks. + RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR +}; + +#if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED) +namespace detail +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(RadixRankAlgorithm algo) noexcept +{ + switch (algo) + { + case RADIX_RANK_BASIC: + return "RADIX_RANK_BASIC"; + case RADIX_RANK_MEMOIZE: + return "RADIX_RANK_MEMOIZE"; + case RADIX_RANK_MATCH: + return "RADIX_RANK_MATCH"; + case RADIX_RANK_MATCH_EARLY_COUNTS_ANY: + return "RADIX_RANK_MATCH_EARLY_COUNTS_ANY"; + case RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR: + return "RADIX_RANK_MATCH_EARLY_COUNTS_ATOMIC_OR"; + } + return ""; +} +} // namespace detail + +inline ::std::ostream& operator<<(::std::ostream& os, RadixRankAlgorithm algo) +{ + return os << CUB_NS_QUALIFIER::detail::to_string(algo); +} +#endif // _CCCL_HOSTED() && !_CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END + +#if __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) +template <::cuda::std::same_as CharT> +struct std::formatter : formatter +{ + template + auto format(const CUB_NS_QUALIFIER::RadixRankAlgorithm& algo, FmtCtx& ctx) const + { + return formatter::format(CUB_NS_QUALIFIER::detail::to_string(algo), ctx); + } +}; +#endif // __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) + +CUB_NAMESPACE_BEGIN + +/** Empty callback implementation */ +template +struct BlockRadixRankEmptyCallback +{ + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(int (&bins)[BINS_PER_THREAD]) {} +}; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document +namespace detail +{ +template +struct warp_in_block_matcher_t +{ + static _CCCL_DEVICE ::cuda::std::uint32_t match_any(::cuda::std::uint32_t label, ::cuda::std::uint32_t warp_id) + { + if (warp_id == static_cast<::cuda::std::uint32_t>(PartialWarpId)) + { + return MatchAny(label); + } + + return MatchAny(label); + } +}; + +template +struct warp_in_block_matcher_t +{ + static _CCCL_DEVICE ::cuda::std::uint32_t match_any(::cuda::std::uint32_t label, ::cuda::std::uint32_t warp_id) + { + return MatchAny(label); + } +}; +} // namespace detail +#endif // _CCCL_DOXYGEN_INVOKED + +//! @rst +//! BlockRadixRank provides operations for ranking unsigned integer types within a CUDA thread block. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - Keys must be in a form suitable for radix ranking (i.e., unsigned bits). +//! - **Important**: BlockRadixRank ranks only ``RadixBits`` bits at a time from the keys, not the entire key. +//! The digit extractor determines which bits are ranked. +//! - @blocked +//! +//! Performance Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - @granularity +//! +//! .. code-block:: c++ +//! +//! #include +//! +//! __global__ void ExampleKernel(...) +//! { +//! constexpr int threads_per_block = 2; +//! constexpr int radix_bits = 5; +//! +//! // Specialize BlockRadixRank for a 1D block of 2 threads +//! using block_radix_rank = cub::BlockRadixRank; +//! using storage_t = typename block_radix_rank::TempStorage; +//! +//! // Allocate shared memory for BlockRadixRank +//! __shared__ storage_t temp_storage; +//! +//! // Obtain a segment of consecutive items that are blocked across threads +//! unsigned int keys[2]; +//! int ranks[2]; +//! ... +//! +//! // Extract the lowest radix_bits from each key +//! cub::BFEDigitExtractor extractor(0, radix_bits); +//! block_radix_rank(temp_storage).RankKeys(keys, ranks, extractor); +//! +//! ... +//! } +//! +//! Suppose the set of input ``keys`` across the block of threads is ``{ [16,10], [9,11] }``. +//! The extractor will rank only the lowest 5 bits: ``{ [16,10], [9,11] }`` (bits 0-4). +//! The corresponding output ``ranks`` in those threads will be ``{ [3,1], [0,2] }``. +//! +//! Re-using dynamically allocating shared memory +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of dynamically shared memory with +//! BlockReduce and how to re-purpose the same memory region. +//! This example can be easily adapted to the storage required by BlockRadixRank. +//! +//! @endrst +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam RadixBits +//! The number of radix bits per digit place +//! +//! @tparam IsDescending +//! Whether or not the sorted-order is high-to-low +//! +//! @tparam MemoizeOuterScan +//! **[optional]** Whether or not to buffer outer raking scan +//! partials to incur fewer shared memory reads at the expense of higher register pressure +//! (default: true for architectures SM35 and newer, false otherwise). +//! See `BlockScanAlgorithm::BLOCK_SCAN_RAKING_MEMOIZE` for more details. +//! +//! @tparam InnerScanAlgorithm +//! **[optional]** The cub::BlockScanAlgorithm algorithm to use (default: cub::BLOCK_SCAN_WARP_SCANS) +//! +//! @tparam SMemConfig +//! **[optional]** Shared memory bank mode (default: `cudaSharedMemBankSizeFourByte`) +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockRadixRank +{ +private: + // Integer type for digit counters (to be packed into words of type PackedCounters) + using DigitCounter = unsigned short; + + // Integer type for packing DigitCounters into columns of shared memory banks + using PackedCounter = + ::cuda::std::_If; + + static constexpr DigitCounter max_tile_size = ::cuda::std::numeric_limits::max(); + + // The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + static constexpr int RADIX_DIGITS = 1 << RadixBits; + + static constexpr int LOG_WARP_THREADS = detail::log2_warp_threads; + static constexpr int WARP_THREADS = 1 << LOG_WARP_THREADS; + static constexpr int WARPS = (BLOCK_THREADS + WARP_THREADS - 1) / WARP_THREADS; + + static constexpr int BYTES_PER_COUNTER = sizeof(DigitCounter); + static constexpr int LOG_BYTES_PER_COUNTER = Log2::VALUE; + + static constexpr int PACKING_RATIO = static_cast(sizeof(PackedCounter) / sizeof(DigitCounter)); + static constexpr int LOG_PACKING_RATIO = Log2::VALUE; + + // Always at least one lane + static constexpr int LOG_COUNTER_LANES = ::cuda::std::max(RadixBits - LOG_PACKING_RATIO, 0); + static constexpr int COUNTER_LANES = 1 << LOG_COUNTER_LANES; + + // The number of packed counters per thread (plus one for padding) + static constexpr int PADDED_COUNTER_LANES = COUNTER_LANES + 1; + static constexpr int RAKING_SEGMENT = PADDED_COUNTER_LANES; + +public: + /// Number of bin-starting offsets tracked per thread + static constexpr int BINS_TRACKED_PER_THREAD = + ::cuda::std::max(1, (RADIX_DIGITS + BLOCK_THREADS - 1) / BLOCK_THREADS); + +private: + /// BlockScan type + using BlockScan = BlockScan; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + struct __align__(16) _TempStorage + { + union Aliasable + { + DigitCounter digit_counters[PADDED_COUNTER_LANES][BLOCK_THREADS][PACKING_RATIO]; + PackedCounter raking_grid[BLOCK_THREADS][RAKING_SEGMENT]; + + } aliasable; + + // Storage for scanning local ranks + typename BlockScan::TempStorage block_scan; + }; +#endif // !_CCCL_DOXYGEN_INVOKED + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + unsigned int linear_tid; + + /// Copy of raking segment, promoted to registers + PackedCounter cached_segment[RAKING_SEGMENT]; + + /** + * Internal storage allocator + */ + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + /** + * Performs upsweep raking reduction, returning the aggregate + */ + _CCCL_DEVICE _CCCL_FORCEINLINE PackedCounter Upsweep() + { + auto& smem_raking_ptr = temp_storage.aliasable.raking_grid[linear_tid]; + if constexpr (MemoizeOuterScan) + { + // Copy data into registers + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < RAKING_SEGMENT; i++) + { + cached_segment[i] = smem_raking_ptr[i]; + } + return cub::ThreadReduce(::cuda::std::span{cached_segment}, ::cuda::std::plus<>{}); + } + else + { + return cub::ThreadReduce(smem_raking_ptr, ::cuda::std::plus<>{}); + } + } + + /// Performs exclusive downsweep raking scan + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveDownsweep(PackedCounter raking_partial) + { + PackedCounter* smem_raking_ptr = temp_storage.aliasable.raking_grid[linear_tid]; + + PackedCounter* raking_ptr = (MemoizeOuterScan) ? cached_segment : smem_raking_ptr; + + // Exclusive raking downsweep scan + detail::ThreadScanExclusive(raking_ptr, raking_ptr, ::cuda::std::plus<>{}, raking_partial); + + if (MemoizeOuterScan) + { + // Copy data back to smem + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < RAKING_SEGMENT; i++) + { + smem_raking_ptr[i] = cached_segment[i]; + } + } + } + + /** + * Reset shared memory digit counters + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ResetCounters() + { + // Reset shared memory digit counters + _CCCL_PRAGMA_UNROLL_FULL() + for (int LANE = 0; LANE < PADDED_COUNTER_LANES; LANE++) + { + *((PackedCounter*) temp_storage.aliasable.digit_counters[LANE][linear_tid]) = 0; + } + } + + /** + * Block-scan prefix callback + */ + struct PrefixCallBack + { + _CCCL_DEVICE _CCCL_FORCEINLINE PackedCounter operator()(PackedCounter block_aggregate) + { + PackedCounter block_prefix = 0; + + // Propagate totals in packed fields + _CCCL_PRAGMA_UNROLL_FULL() + for (int PACKED = 1; PACKED < PACKING_RATIO; PACKED++) + { + block_prefix += block_aggregate << (sizeof(DigitCounter) * 8 * PACKED); + } + + return block_prefix; + } + }; + + /** + * Scan shared memory digit counters. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void ScanCounters() + { + // Upsweep scan + PackedCounter raking_partial = Upsweep(); + + // Compute exclusive sum + PackedCounter exclusive_partial; + PrefixCallBack prefix_call_back; + BlockScan(temp_storage.block_scan).ExclusiveSum(raking_partial, exclusive_partial, prefix_call_back); + + // Downsweep scan with exclusive partial + ExclusiveDownsweep(exclusive_partial); + } + +public: + /// @smemstorage{BlockScan} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using a private static allocation of shared memory as temporary storage. + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRadixRank() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @brief Collective constructor using the specified memory allocation as temporary storage. + * + * @param[in] temp_storage + * Reference to memory allocation having layout type TempStorage + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRadixRank(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Raking + //! @{ + + /** + * @brief Rank keys. + * + * @param[in] keys + * Keys for this tile + * + * @param[out] ranks + * For each key, the local rank within the tile + * + * @param[in] digit_extractor + * The digit extractor + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(UnsignedBits (&keys)[KEYS_PER_THREAD], int (&ranks)[KEYS_PER_THREAD], DigitExtractorT digit_extractor) + { + static_assert(BLOCK_THREADS * KEYS_PER_THREAD <= max_tile_size, + "DigitCounter type is too small to hold this number of keys"); + + DigitCounter thread_prefixes[KEYS_PER_THREAD]; // For each key, the count of previous keys in this tile having the + // same digit + DigitCounter* digit_counters[KEYS_PER_THREAD]; // For each key, the byte-offset of its corresponding digit counter + // in smem + + // Reset shared memory digit counters + ResetCounters(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < KEYS_PER_THREAD; ++ITEM) + { + // Get digit + ::cuda::std::uint32_t digit = digit_extractor.Digit(keys[ITEM]); + + // Get sub-counter + ::cuda::std::uint32_t sub_counter = digit >> LOG_COUNTER_LANES; + + // Get counter lane + ::cuda::std::uint32_t counter_lane = digit & (COUNTER_LANES - 1); + + if (IsDescending) + { + sub_counter = PACKING_RATIO - 1 - sub_counter; + counter_lane = COUNTER_LANES - 1 - counter_lane; + } + + // Pointer to smem digit counter + digit_counters[ITEM] = &temp_storage.aliasable.digit_counters[counter_lane][linear_tid][sub_counter]; + + // Load thread-exclusive prefix + thread_prefixes[ITEM] = *digit_counters[ITEM]; + + // Store inclusive prefix + *digit_counters[ITEM] = thread_prefixes[ITEM] + 1; + } + + __syncthreads(); + + // Scan shared memory counters + ScanCounters(); + + __syncthreads(); + + // Extract the local ranks of each key + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < KEYS_PER_THREAD; ++ITEM) + { + // Add in thread block exclusive prefix + ranks[ITEM] = thread_prefixes[ITEM] + *digit_counters[ITEM]; + } + } + + /** + * @brief Rank keys. For the lower @p RADIX_DIGITS threads, digit counts for each digit are + * provided for the corresponding thread. + * + * @param[in] keys + * Keys for this tile + * + * @param[out] ranks + * For each key, the local rank within the tile (out parameter) + * + * @param[in] digit_extractor + * The digit extractor + * + * @param[out] exclusive_digit_prefix + * The exclusive prefix sum for the digits + * [(threadIdx.x * BINS_TRACKED_PER_THREAD) + * ... + * (threadIdx.x * BINS_TRACKED_PER_THREAD) + BINS_TRACKED_PER_THREAD - 1] + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(UnsignedBits (&keys)[KEYS_PER_THREAD], + int (&ranks)[KEYS_PER_THREAD], + DigitExtractorT digit_extractor, + int (&exclusive_digit_prefix)[BINS_TRACKED_PER_THREAD]) + { + static_assert(BLOCK_THREADS * KEYS_PER_THREAD <= max_tile_size, + "DigitCounter type is too small to hold this number of keys"); + + // Rank keys + RankKeys(keys, ranks, digit_extractor); + + // Get the inclusive and exclusive digit totals corresponding to the calling thread. + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (linear_tid * BINS_TRACKED_PER_THREAD) + track; + + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + if (IsDescending) + { + bin_idx = RADIX_DIGITS - bin_idx - 1; + } + + // Obtain ex/inclusive digit counts. (Unfortunately these all reside in the + // first counter column, resulting in unavoidable bank conflicts.) + unsigned int counter_lane = (bin_idx & (COUNTER_LANES - 1)); + unsigned int sub_counter = bin_idx >> (LOG_COUNTER_LANES); + + exclusive_digit_prefix[track] = temp_storage.aliasable.digit_counters[counter_lane][0][sub_counter]; + } + } + } + + //! @} +}; + +/** + * Radix-rank using match.any + */ +template +class BlockRadixRankMatch +{ +private: + using RankT = int32_t; + using DigitCounterT = int32_t; + + // The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + static constexpr int RADIX_DIGITS = 1 << RadixBits; + + static constexpr int LOG_WARP_THREADS = detail::log2_warp_threads; + static constexpr int WARP_THREADS = 1 << LOG_WARP_THREADS; + static constexpr int PARTIAL_WARP_THREADS = BLOCK_THREADS % WARP_THREADS; + static constexpr int WARPS = (BLOCK_THREADS + WARP_THREADS - 1) / WARP_THREADS; + + static constexpr int PADDED_WARPS = ((WARPS & 0x1) == 0) ? WARPS + 1 : WARPS; + + static constexpr int COUNTERS = PADDED_WARPS * RADIX_DIGITS; + static constexpr int RAKING_SEGMENT = (COUNTERS + BLOCK_THREADS - 1) / BLOCK_THREADS; + static constexpr int PADDED_RAKING_SEGMENT = ((RAKING_SEGMENT & 0x1) == 0) ? RAKING_SEGMENT + 1 : RAKING_SEGMENT; + +public: + /// Number of bin-starting offsets tracked per thread + static constexpr int BINS_TRACKED_PER_THREAD = + ::cuda::std::max(1, (RADIX_DIGITS + BLOCK_THREADS - 1) / BLOCK_THREADS); + +private: + /// BlockScan type + using BlockScanT = BlockScan; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + struct __align__(16) _TempStorage + { + typename BlockScanT::TempStorage block_scan; + + union __align__(16) Aliasable + { + volatile DigitCounterT warp_digit_counters[RADIX_DIGITS][PADDED_WARPS]; + DigitCounterT raking_grid[BLOCK_THREADS][PADDED_RAKING_SEGMENT]; + } aliasable; + }; +#endif // !_CCCL_DOXYGEN_INVOKED + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + unsigned int linear_tid; + +public: + /// @smemstorage{BlockRadixRankMatch} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + /** + * @brief Collective constructor using the specified memory allocation as temporary storage. + * + * @param[in] temp_storage + * Reference to memory allocation having layout type TempStorage + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRadixRankMatch(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Raking + //! @{ + + /** + * @brief Computes the count of keys for each digit value, and calls the + * callback with the array of key counts. + * + * @tparam CountsCallback The callback type. It should implement an instance + * overload of operator()(int (&bins)[BINS_TRACKED_PER_THREAD]), where bins + * is an array of key counts for each digit value distributed in block + * distribution among the threads of the thread block. Key counts can be + * used, to update other data structures in global or shared + * memory. Depending on the implementation of the ranking algoirhtm + * (see BlockRadixRankMatchEarlyCounts), key counts may become available + * early, therefore, they are returned through a callback rather than a + * separate output parameter of RankKeys(). + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void CallBack(CountsCallback callback) + { + int bins[BINS_TRACKED_PER_THREAD]; + // Get count for each digit + + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (linear_tid * BINS_TRACKED_PER_THREAD) + track; + constexpr int TILE_ITEMS = KEYS_PER_THREAD * BLOCK_THREADS; + + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + if (IsDescending) + { + bin_idx = RADIX_DIGITS - bin_idx - 1; + bins[track] = (bin_idx > 0 ? temp_storage.aliasable.warp_digit_counters[bin_idx - 1][0] : TILE_ITEMS) + - temp_storage.aliasable.warp_digit_counters[bin_idx][0]; + } + else + { + bins[track] = + (bin_idx < RADIX_DIGITS - 1 ? temp_storage.aliasable.warp_digit_counters[bin_idx + 1][0] : TILE_ITEMS) + - temp_storage.aliasable.warp_digit_counters[bin_idx][0]; + } + } + } + callback(bins); + } + + /** + * @brief Rank keys. + * + * @param[in] keys + * Keys for this tile + * + * @param[out] ranks + * For each key, the local rank within the tile + * + * @param[in] digit_extractor + * The digit extractor + * + * @param[in] callback + * Callback to receive digit counts + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(UnsignedBits (&keys)[KEYS_PER_THREAD], + int (&ranks)[KEYS_PER_THREAD], + DigitExtractorT digit_extractor, + CountsCallback callback) + { + // Initialize shared digit counters + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < PADDED_RAKING_SEGMENT; ++ITEM) + { + temp_storage.aliasable.raking_grid[linear_tid][ITEM] = 0; + } + + __syncthreads(); + + // Each warp will strip-mine its section of input, one strip at a time + + volatile DigitCounterT* digit_counters[KEYS_PER_THREAD]; + ::cuda::std::uint32_t warp_id = linear_tid >> LOG_WARP_THREADS; + ::cuda::std::uint32_t lane_mask_lt = ::cuda::ptx::get_sreg_lanemask_lt(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < KEYS_PER_THREAD; ++ITEM) + { + // My digit + ::cuda::std::uint32_t digit = digit_extractor.Digit(keys[ITEM]); + + if (IsDescending) + { + digit = RADIX_DIGITS - digit - 1; + } + + // Mask of peers who have same digit as me + ::cuda::std::uint32_t peer_mask = + detail::warp_in_block_matcher_t::match_any(digit, warp_id); + + // Pointer to smem digit counter for this key + digit_counters[ITEM] = &temp_storage.aliasable.warp_digit_counters[digit][warp_id]; + + // Number of occurrences in previous strips + DigitCounterT warp_digit_prefix = *digit_counters[ITEM]; + + // Warp-sync + __syncwarp(0xFFFFFFFF); + + // Number of peers having same digit as me + int32_t digit_count = ::cuda::std::popcount(peer_mask); + + // Number of lower-ranked peers having same digit seen so far + int32_t peer_digit_prefix = ::cuda::std::popcount(peer_mask & lane_mask_lt); + + if (peer_digit_prefix == 0) + { + // First thread for each digit updates the shared warp counter + *digit_counters[ITEM] = DigitCounterT(warp_digit_prefix + digit_count); + } + + // Warp-sync + __syncwarp(0xFFFFFFFF); + + // Number of prior keys having same digit + ranks[ITEM] = warp_digit_prefix + DigitCounterT(peer_digit_prefix); + } + + __syncthreads(); + + // Scan warp counters + + DigitCounterT scan_counters[PADDED_RAKING_SEGMENT]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < PADDED_RAKING_SEGMENT; ++ITEM) + { + scan_counters[ITEM] = temp_storage.aliasable.raking_grid[linear_tid][ITEM]; + } + + BlockScanT(temp_storage.block_scan).ExclusiveSum(scan_counters, scan_counters); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < PADDED_RAKING_SEGMENT; ++ITEM) + { + temp_storage.aliasable.raking_grid[linear_tid][ITEM] = scan_counters[ITEM]; + } + + __syncthreads(); + if (!::cuda::std::is_same_v>) + { + CallBack(callback); + } + + // Seed ranks with counter values from previous warps + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < KEYS_PER_THREAD; ++ITEM) + { + ranks[ITEM] += *digit_counters[ITEM]; + } + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(UnsignedBits (&keys)[KEYS_PER_THREAD], int (&ranks)[KEYS_PER_THREAD], DigitExtractorT digit_extractor) + { + RankKeys(keys, ranks, digit_extractor, BlockRadixRankEmptyCallback()); + } + + /** + * @brief Rank keys. For the lower @p RADIX_DIGITS threads, digit counts for each digit are + * provided for the corresponding thread. + * + * @param[in] keys + * Keys for this tile + * + * @param[out] ranks + * For each key, the local rank within the tile (out parameter) + * + * @param[in] digit_extractor + * The digit extractor + * + * @param[out] exclusive_digit_prefix + * The exclusive prefix sum for the digits + * [(threadIdx.x * BINS_TRACKED_PER_THREAD) + * ... + * (threadIdx.x * BINS_TRACKED_PER_THREAD) + BINS_TRACKED_PER_THREAD - 1] + * + * @param[in] callback + * Callback to receive digit counts + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void RankKeys( + UnsignedBits (&keys)[KEYS_PER_THREAD], + int (&ranks)[KEYS_PER_THREAD], + DigitExtractorT digit_extractor, + int (&exclusive_digit_prefix)[BINS_TRACKED_PER_THREAD], + CountsCallback callback) + { + RankKeys(keys, ranks, digit_extractor, callback); + + // Get exclusive count for each digit + _CCCL_PRAGMA_UNROLL_FULL() + for (int track = 0; track < BINS_TRACKED_PER_THREAD; ++track) + { + int bin_idx = (linear_tid * BINS_TRACKED_PER_THREAD) + track; + + if ((BLOCK_THREADS == RADIX_DIGITS) || (bin_idx < RADIX_DIGITS)) + { + if (IsDescending) + { + bin_idx = RADIX_DIGITS - bin_idx - 1; + } + + exclusive_digit_prefix[track] = temp_storage.aliasable.warp_digit_counters[bin_idx][0]; + } + } + } + + /** + * @param[in] keys + * Keys for this tile + * + * @param[out] ranks + * For each key, the local rank within the tile (out parameter) + * + * @param[out] exclusive_digit_prefix + * The exclusive prefix sum for the digits + * [(threadIdx.x * BINS_TRACKED_PER_THREAD) + * ... + * (threadIdx.x * BINS_TRACKED_PER_THREAD) + BINS_TRACKED_PER_THREAD - 1] + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(UnsignedBits (&keys)[KEYS_PER_THREAD], + int (&ranks)[KEYS_PER_THREAD], + DigitExtractorT digit_extractor, + int (&exclusive_digit_prefix)[BINS_TRACKED_PER_THREAD]) + { + RankKeys( + keys, ranks, digit_extractor, exclusive_digit_prefix, BlockRadixRankEmptyCallback()); + } + + //! @} +}; + +enum WarpMatchAlgorithm +{ + WARP_MATCH_ANY, + WARP_MATCH_ATOMIC_OR +}; + +/** + * Radix-rank using matching which computes the counts of keys for each digit + * value early, at the expense of doing more work. This may be useful e.g. for + * decoupled look-back, where it reduces the time other thread blocks need to + * wait for digit counts to become available. + */ +template +struct BlockRadixRankMatchEarlyCounts +{ + // constants + static constexpr int BLOCK_THREADS = BlockDimX; + static constexpr int RADIX_DIGITS = 1 << RadixBits; + static constexpr int BINS_PER_THREAD = (RADIX_DIGITS + BLOCK_THREADS - 1) / BLOCK_THREADS; + static constexpr int BINS_TRACKED_PER_THREAD = BINS_PER_THREAD; + static constexpr int FULL_BINS = BINS_PER_THREAD * BLOCK_THREADS == RADIX_DIGITS; + static constexpr int WARP_THREADS = detail::warp_threads; + static constexpr int PARTIAL_WARP_THREADS = BLOCK_THREADS % WARP_THREADS; + static constexpr int BLOCK_WARPS = BLOCK_THREADS / WARP_THREADS; + static constexpr int PARTIAL_WARP_ID = BLOCK_WARPS - 1; + static constexpr int WARP_MASK = ~0; + static constexpr int NUM_MATCH_MASKS = MATCH_ALGORITHM == WARP_MATCH_ATOMIC_OR ? BLOCK_WARPS : 0; + // Guard against declaring zero-sized array: + static constexpr int MATCH_MASKS_ALLOC_SIZE = NUM_MATCH_MASKS < 1 ? 1 : NUM_MATCH_MASKS; + + // types + using BlockScan = cub::BlockScan; + + struct TempStorage + { + union + { + int warp_offsets[BLOCK_WARPS][RADIX_DIGITS]; + int warp_histograms[BLOCK_WARPS][RADIX_DIGITS][NUM_PARTS]; + }; + + ::cuda::std::uint32_t match_masks[MATCH_MASKS_ALLOC_SIZE][RADIX_DIGITS]; + + typename BlockScan::TempStorage prefix_tmp; + }; + + TempStorage& temp_storage; + + // internal ranking implementation + template + struct BlockRadixRankMatchInternal + { + TempStorage& s; + DigitExtractorT digit_extractor; + CountsCallback callback; + int warp; + int lane; + + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t Digit(UnsignedBits key) + { + ::cuda::std::uint32_t digit = digit_extractor.Digit(key); + return IsDescending ? RADIX_DIGITS - 1 - digit : digit; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE int ThreadBin(int u) + { + int bin = threadIdx.x * BINS_PER_THREAD + u; + return IsDescending ? RADIX_DIGITS - 1 - bin : bin; + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ComputeHistogramsWarp(UnsignedBits (&keys)[KEYS_PER_THREAD]) + { + // int* warp_offsets = &s.warp_offsets[warp][0]; + int (&warp_histograms)[RADIX_DIGITS][NUM_PARTS] = s.warp_histograms[warp]; + + // compute warp-private histograms + _CCCL_PRAGMA_UNROLL_FULL() + for (int bin = lane; bin < RADIX_DIGITS; bin += WARP_THREADS) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int part = 0; part < NUM_PARTS; ++part) + { + warp_histograms[bin][part] = 0; + } + } + if constexpr (MATCH_ALGORITHM == WARP_MATCH_ATOMIC_OR) + { + ::cuda::std::uint32_t* match_masks = &s.match_masks[warp][0]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int bin = lane; bin < RADIX_DIGITS; bin += WARP_THREADS) + { + match_masks[bin] = 0; + } + } + __syncwarp(WARP_MASK); + + // compute private per-part histograms + int part = lane % NUM_PARTS; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < KEYS_PER_THREAD; ++u) + { + atomicAdd(&warp_histograms[Digit(keys[u])][part], 1); + } + + // sum different parts; + // no extra work is necessary if NUM_PARTS == 1 + if constexpr (NUM_PARTS > 1) + { + __syncwarp(WARP_MASK); + // TODO: handle RADIX_DIGITS % WARP_THREADS != 0 if it becomes necessary + constexpr int WARP_BINS_PER_THREAD = RADIX_DIGITS / WARP_THREADS; + int bins[WARP_BINS_PER_THREAD]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < WARP_BINS_PER_THREAD; ++u) + { + int bin = lane + u * WARP_THREADS; + bins[u] = cub::ThreadReduce(warp_histograms[bin], ::cuda::std::plus<>{}); + } + __syncthreads(); + + // store the resulting histogram in shared memory + int* warp_offsets = &s.warp_offsets[warp][0]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < WARP_BINS_PER_THREAD; ++u) + { + int bin = lane + u * WARP_THREADS; + warp_offsets[bin] = bins[u]; + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ComputeOffsetsWarpUpsweep(int (&bins)[BINS_PER_THREAD]) + { + // sum up warp-private histograms + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + bins[u] = 0; + int bin = ThreadBin(u); + if (FULL_BINS || (bin >= 0 && bin < RADIX_DIGITS)) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int j_warp = 0; j_warp < BLOCK_WARPS; ++j_warp) + { + int warp_offset = s.warp_offsets[j_warp][bin]; + s.warp_offsets[j_warp][bin] = bins[u]; + bins[u] += warp_offset; + } + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ComputeOffsetsWarpDownsweep(int (&offsets)[BINS_PER_THREAD]) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < BINS_PER_THREAD; ++u) + { + int bin = ThreadBin(u); + if (FULL_BINS || (bin >= 0 && bin < RADIX_DIGITS)) + { + int digit_offset = offsets[u]; + _CCCL_PRAGMA_UNROLL_FULL() + for (int j_warp = 0; j_warp < BLOCK_WARPS; ++j_warp) + { + s.warp_offsets[j_warp][bin] += digit_offset; + } + } + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ComputeRanksItem( + UnsignedBits (&keys)[KEYS_PER_THREAD], int (&ranks)[KEYS_PER_THREAD], detail::constant_t) + { + // compute key ranks + ::cuda::std::uint32_t lane_mask = 1u << lane; + int* warp_offsets = &s.warp_offsets[warp][0]; + ::cuda::std::uint32_t* match_masks = &s.match_masks[warp][0]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < KEYS_PER_THREAD; ++u) + { + ::cuda::std::uint32_t bin = Digit(keys[u]); + ::cuda::std::uint32_t* p_match_mask = &match_masks[bin]; + atomicOr(p_match_mask, lane_mask); + __syncwarp(WARP_MASK); + ::cuda::std::uint32_t bin_mask = *p_match_mask; + // TODO(bgruber): __bit_log2 regresses cub.bench.radix_sort.keys.base up to 30% on H200, see cccl_private/#586 + // int leader = ::cuda::std::__bit_log2(bin_mask); + int leader = (WARP_THREADS - 1) - ::cuda::std::countl_zero(bin_mask); + int warp_offset = 0; + int popc = ::cuda::std::popcount(bin_mask & ::cuda::ptx::get_sreg_lanemask_le()); + if (lane == leader) + { + // atomic is a bit faster + warp_offset = atomicAdd(&warp_offsets[bin], popc); + } + warp_offset = __shfl_sync(WARP_MASK, warp_offset, leader); + if (lane == leader) + { + *p_match_mask = 0; + } + __syncwarp(WARP_MASK); + ranks[u] = warp_offset + popc - 1; + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void ComputeRanksItem( + UnsignedBits (&keys)[KEYS_PER_THREAD], int (&ranks)[KEYS_PER_THREAD], detail::constant_t) + { + // compute key ranks + int* warp_offsets = &s.warp_offsets[warp][0]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int u = 0; u < KEYS_PER_THREAD; ++u) + { + ::cuda::std::uint32_t bin = Digit(keys[u]); + ::cuda::std::uint32_t bin_mask = + detail::warp_in_block_matcher_t::match_any(bin, warp); + // TODO(bgruber): __bit_log2 regresses cub.bench.radix_sort.keys.base up to 30% on H200, see cccl_private/#586 + // int leader = ::cuda::std::__bit_log2(bin_mask); + int leader = (WARP_THREADS - 1) - ::cuda::std::countl_zero(bin_mask); + int warp_offset = 0; + int popc = ::cuda::std::popcount(bin_mask & ::cuda::ptx::get_sreg_lanemask_le()); + if (lane == leader) + { + // atomic is a bit faster + warp_offset = atomicAdd(&warp_offsets[bin], popc); + } + warp_offset = __shfl_sync(WARP_MASK, warp_offset, leader); + ranks[u] = warp_offset + popc - 1; + } + } + + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(UnsignedBits (&keys)[KEYS_PER_THREAD], + int (&ranks)[KEYS_PER_THREAD], + int (&exclusive_digit_prefix)[BINS_PER_THREAD]) + { + ComputeHistogramsWarp(keys); + + __syncthreads(); + int bins[BINS_PER_THREAD]; + ComputeOffsetsWarpUpsweep(bins); + callback(bins); + + BlockScan(s.prefix_tmp).ExclusiveSum(bins, exclusive_digit_prefix); + + ComputeOffsetsWarpDownsweep(exclusive_digit_prefix); + __syncthreads(); + ComputeRanksItem(keys, ranks, detail::constant_v); + } + + _CCCL_DEVICE _CCCL_FORCEINLINE + BlockRadixRankMatchInternal(TempStorage& temp_storage, DigitExtractorT digit_extractor, CountsCallback callback) + : s(temp_storage) + , digit_extractor(digit_extractor) + , callback(callback) + , warp(static_cast(threadIdx.x / WARP_THREADS)) + , lane(static_cast(::cuda::ptx::get_sreg_laneid())) + {} + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRadixRankMatchEarlyCounts(TempStorage& temp_storage) + : temp_storage(temp_storage) + {} + + /** + * @brief Rank keys. For the lower @p RADIX_DIGITS threads, digit counts for each digit are + * provided for the corresponding thread. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void RankKeys( + UnsignedBits (&keys)[KEYS_PER_THREAD], + int (&ranks)[KEYS_PER_THREAD], + DigitExtractorT digit_extractor, + int (&exclusive_digit_prefix)[BINS_PER_THREAD], + CountsCallback callback) + { + BlockRadixRankMatchInternal internal( + temp_storage, digit_extractor, callback); + internal.RankKeys(keys, ranks, exclusive_digit_prefix); + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(UnsignedBits (&keys)[KEYS_PER_THREAD], + int (&ranks)[KEYS_PER_THREAD], + DigitExtractorT digit_extractor, + int (&exclusive_digit_prefix)[BINS_PER_THREAD]) + { + using CountsCallback = BlockRadixRankEmptyCallback; + BlockRadixRankMatchInternal internal( + temp_storage, digit_extractor, CountsCallback()); + internal.RankKeys(keys, ranks, exclusive_digit_prefix); + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(UnsignedBits (&keys)[KEYS_PER_THREAD], int (&ranks)[KEYS_PER_THREAD], DigitExtractorT digit_extractor) + { + int exclusive_digit_prefix[BINS_PER_THREAD]; + RankKeys(keys, ranks, digit_extractor, exclusive_digit_prefix); + } +}; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document +namespace detail +{ +// `BlockRadixRank` doesn't conform to the typical pattern, not exposing the algorithm +// template parameter. Other algorithms don't provide the same template parameters, not allowing +// multi-dimensional thread block specializations. +// +// TODO(senior-zero) for 3.0: +// - Put existing implementations into the detail namespace +// - Support multi-dimensional thread blocks in the rest of implementations +// - Repurpose BlockRadixRank as an entry name with the algorithm template parameter +template +using block_radix_rank_t = ::cuda::std::_If< + RankAlgorithm == RADIX_RANK_BASIC, + BlockRadixRank, + ::cuda::std::_If< + RankAlgorithm == RADIX_RANK_MEMOIZE, + BlockRadixRank, + ::cuda::std::_If< + RankAlgorithm == RADIX_RANK_MATCH, + BlockRadixRankMatch, + ::cuda::std::_If< + RankAlgorithm == RADIX_RANK_MATCH_EARLY_COUNTS_ANY, + BlockRadixRankMatchEarlyCounts, + BlockRadixRankMatchEarlyCounts>>>>; +} // namespace detail +#endif // _CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_radix_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_radix_sort.cuh new file mode 100644 index 00000000..f50ae190 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_radix_sort.cuh @@ -0,0 +1,2191 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * The cub::BlockRadixSort class provides [collective](../index.html#sec0) methods for radix + * sorting of items partitioned across a CUDA thread block. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! BlockRadixSort class provides :ref:`collective ` methods for sorting +//! items partitioned across a CUDA thread block using a radix sorting method. +//! +//! .. image:: ../../img/sorting_logo.png +//! :align: center +//! +//! Overview +//! -------------------------------------------------- +//! +//! The `radix sorting method `_ arranges +//! items into ascending order. It relies upon a positional representation for +//! keys, i.e., each key is comprised of an ordered sequence of symbols (e.g., digits, +//! characters, etc.) specified from least-significant to most-significant. For a +//! given input sequence of keys and a set of rules specifying a total ordering +//! of the symbolic alphabet, the radix sorting method produces a lexicographic +//! ordering of those keys. +//! +//! @rowmajor +//! +//! Supported Types +//! -------------------------------------------------- +//! +//! BlockRadixSort can sort all of the built-in C++ numeric primitive types +//! (``unsigned char``, ``int``, ``double``, etc.) as well as CUDA's ``__half`` +//! half-precision floating-point type. User-defined types are supported as long +//! as decomposer object is provided. +//! +//! Floating-Point Special Cases +//! -------------------------------------------------- +//! +//! - Positive and negative zeros are considered equivalent, and will be treated +//! as such in the output. +//! - No special handling is implemented for NaN values; these are sorted +//! according to their bit representations after any transformations. +//! +//! Bitwise Key Transformations +//! -------------------------------------------------- +//! +//! Although the direct radix sorting method can only be applied to unsigned +//! integral types, BlockRadixSort is able to sort signed and floating-point +//! types via simple bit-wise transformations that ensure lexicographic key +//! ordering. +//! +//! These transformations must be considered when restricting the +//! ``[begin_bit, end_bit)`` range, as the bitwise transformations will occur +//! before the bit-range truncation. +//! +//! Any transformations applied to the keys prior to sorting are reversed +//! while writing to the final output buffer. +//! +//! Type Specific Bitwise Transformations +//! -------------------------------------------------- +//! +//! To convert the input values into a radix-sortable bitwise representation, +//! the following transformations take place prior to sorting: +//! +//! * For unsigned integral values, the keys are used directly. +//! * For signed integral values, the sign bit is inverted. +//! * For positive floating point values, the sign bit is inverted. +//! * For negative floating point values, the full key is inverted. +//! +//! No Descending Sort Transformations +//! -------------------------------------------------- +//! +//! Unlike ``DeviceRadixSort``, ``BlockRadixSort`` does not invert the input key bits +//! when performing a descending sort. Instead, it has special logic to reverse +//! the order of the keys while sorting. +//! +//! Stability +//! -------------------------------------------------- +//! +//! BlockRadixSort is stable. For floating-point types -0.0 and +0.0 +//! are considered equal and appear in the result in the same order as they +//! appear in the input. +//! +//! +//! Performance Considerations +//! -------------------------------------------------- +//! +//! * @granularity +//! +//! A Simple Example +//! -------------------------------------------------- +//! +//! @blockcollective{BlockRadixSort} +//! +//! The code snippet below illustrates a sort of 512 integer keys that +//! are partitioned in a [blocked arrangement](../index.html#sec5sec3) across 128 threads +//! where each thread owns 4 consecutive items. +//! +//! .. code-block:: c++ +//! +//! #include // or equivalently +//! +//! __global__ void kernel(...) +//! { +//! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer items each +//! using BlockRadixSort = cub::BlockRadixSort; +//! +//! // Allocate shared memory for BlockRadixSort +//! __shared__ typename BlockRadixSort::TempStorage temp_storage; +//! +//! // Obtain a segment of consecutive items that are blocked across threads +//! int thread_keys[4]; +//! ... +//! +//! // Collectively sort the keys +//! BlockRadixSort(temp_storage).Sort(thread_keys); +//! +//! ... +//! } +//! +//! Suppose the set of input ``thread_keys`` across the block of threads is +//! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. +//! The corresponding output ``thread_keys`` in those threads will be +//! ``{ [0,1,2,3], [4,5,6,7], [8,9,10,11], ..., [508,509,510,511] }``. +//! +//! Re-using dynamically allocating shared memory +//! -------------------------------------------------- +//! +//! The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of dynamically shared memory with +//! BlockReduce and how to re-purpose the same memory region. +//! +//! This example can be easily adapted to the storage required by BlockRadixSort. +//! @endrst +//! +//! @tparam KeyT +//! KeyT type +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam ItemsPerThread +//! The number of items per thread +//! +//! @tparam ValueT +//! **[optional]** ValueT type (default: cub::NullType, which indicates a keys-only sort) +//! +//! @tparam RadixBits +//! **[optional]** The number of radix bits per digit place (default: 4 bits) +//! +//! @tparam MemoizeOuterScan +//! **[optional]** Whether or not to buffer outer raking scan partials to incur fewer shared memory +//! reads at the expense of higher register pressure (default: true for architectures SM35 and +//! newer, false otherwise). +//! +//! @tparam InnerScanAlgorithm +//! **[optional]** The cub::BlockScanAlgorithm algorithm to use +//! (default: cub::BLOCK_SCAN_WARP_SCANS) +//! +//! @tparam SMemConfig +//! **[optional]*8 Shared memory bank mode (default: `cudaSharedMemBankSizeFourByte`) +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockRadixSort +{ +private: + /****************************************************************************** + * Constants and type definitions + ******************************************************************************/ + + // The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + // Whether or not there are values to be trucked along with keys + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + + // KeyT traits and unsigned bits type + using traits = detail::radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion = typename traits::bit_ordered_conversion_policy; + + /// Ascending BlockRadixRank utility type + using AscendingBlockRadixRank = + BlockRadixRank; + + /// Descending BlockRadixRank utility type + using DescendingBlockRadixRank = + BlockRadixRank; + + /// Digit extractor type + using fundamental_digit_extractor_t = BFEDigitExtractor; + + /// BlockExchange utility type for keys + using BlockExchangeKeys = BlockExchange; + + /// BlockExchange utility type for values + using BlockExchangeValues = BlockExchange; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + /// Shared memory storage layout type + union _TempStorage + { + typename AscendingBlockRadixRank::TempStorage asending_ranking_storage; + typename DescendingBlockRadixRank::TempStorage descending_ranking_storage; + typename BlockExchangeKeys::TempStorage exchange_keys; + typename BlockExchangeValues::TempStorage exchange_values; + }; +#endif // _CCCL_DOXYGEN_INVOKED + + /****************************************************************************** + * Thread fields + ******************************************************************************/ + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + unsigned int linear_tid; + + /****************************************************************************** + * Utility methods + ******************************************************************************/ + + /// Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + /// Rank keys (specialized for ascending sort) + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(bit_ordered_type (&unsigned_keys)[ItemsPerThread], + int (&ranks)[ItemsPerThread], + DigitExtractorT digit_extractor, + ::cuda::std::false_type /*is_descending*/) + { + AscendingBlockRadixRank(temp_storage.asending_ranking_storage).RankKeys(unsigned_keys, ranks, digit_extractor); + } + + /// Rank keys (specialized for descending sort) + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + RankKeys(bit_ordered_type (&unsigned_keys)[ItemsPerThread], + int (&ranks)[ItemsPerThread], + DigitExtractorT digit_extractor, + ::cuda::std::true_type /*is_descending*/) + { + DescendingBlockRadixRank(temp_storage.descending_ranking_storage).RankKeys(unsigned_keys, ranks, digit_extractor); + } + + /// ExchangeValues (specialized for key-value sort, to-blocked arrangement) + _CCCL_DEVICE _CCCL_FORCEINLINE void ExchangeValues( + ValueT (&values)[ItemsPerThread], + int (&ranks)[ItemsPerThread], + ::cuda::std::false_type /*is_keys_only*/, + ::cuda::std::true_type /*is_blocked*/) + { + __syncthreads(); + + // Exchange values through shared memory in blocked arrangement + BlockExchangeValues(temp_storage.exchange_values).ScatterToBlocked(values, ranks); + } + + /// ExchangeValues (specialized for key-value sort, to-striped arrangement) + _CCCL_DEVICE _CCCL_FORCEINLINE void ExchangeValues( + ValueT (&values)[ItemsPerThread], + int (&ranks)[ItemsPerThread], + ::cuda::std::false_type /*is_keys_only*/, + ::cuda::std::false_type /*is_blocked*/) + { + __syncthreads(); + + // Exchange values through shared memory in blocked arrangement + BlockExchangeValues(temp_storage.exchange_values).ScatterToStriped(values, ranks); + } + + /// ExchangeValues (specialized for keys-only sort) + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExchangeValues( + ValueT (& /*values*/)[ItemsPerThread], + int (& /*ranks*/)[ItemsPerThread], + ::cuda::std::true_type /*is_keys_only*/, + ::cuda::std::bool_constant /*is_blocked*/) + {} + + /** + * @brief Sort blocked arrangement + * + * @param keys + * Keys to sort + * + * @param values + * Values to sort + * + * @param begin_bit + * The beginning (least-significant) bit index needed for key comparison + * + * @param end_bit + * The past-the-end (most-significant) bit index needed for key comparison + * + * @param is_descending + * Tag whether is a descending-order sort + * + * @param is_keys_only + * Tag whether is keys-only sort + * + * @param decomposer + * Callable object responsible for decomposing a key into a tuple of references to its + * constituent arithmetic types + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SortBlocked( + KeyT (&keys)[ItemsPerThread], + ValueT (&values)[ItemsPerThread], + int begin_bit, + int end_bit, + ::cuda::std::bool_constant is_descending, + ::cuda::std::bool_constant is_keys_only, + DecomposerT decomposer = {}) + { + bit_ordered_type(&unsigned_keys)[ItemsPerThread] = reinterpret_cast(keys); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int KEY = 0; KEY < ItemsPerThread; KEY++) + { + unsigned_keys[KEY] = bit_ordered_conversion::to_bit_ordered(decomposer, unsigned_keys[KEY]); + } + + // Radix sorting passes + while (true) + { + int pass_bits = ::cuda::std::min(RadixBits, end_bit - begin_bit); + auto digit_extractor = + traits::template digit_extractor(begin_bit, pass_bits, decomposer); + + // Rank the blocked keys + int ranks[ItemsPerThread]; + RankKeys(unsigned_keys, ranks, digit_extractor, is_descending); + begin_bit += RadixBits; + + __syncthreads(); + + // Exchange keys through shared memory in blocked arrangement + BlockExchangeKeys(temp_storage.exchange_keys).ScatterToBlocked(keys, ranks); + + // Exchange values through shared memory in blocked arrangement + ExchangeValues(values, ranks, is_keys_only, ::cuda::std::true_type()); + + // Quit if done + if (begin_bit >= end_bit) + { + break; + } + + __syncthreads(); + } + + // Untwiddle bits if necessary + _CCCL_PRAGMA_UNROLL_FULL() + for (int KEY = 0; KEY < ItemsPerThread; KEY++) + { + unsigned_keys[KEY] = bit_ordered_conversion::from_bit_ordered(decomposer, unsigned_keys[KEY]); + } + } + +public: +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + + /** + * @brief Sort blocked -> striped arrangement + * + * @param keys + * Keys to sort + * + * @param values + * Values to sort + * + * @param begin_bit + * The beginning (least-significant) bit index needed for key comparison + * + * @param end_bit + * The past-the-end (most-significant) bit index needed for key comparison + * + * @param is_descending + * Tag whether is a descending-order sort + * + * @param is_keys_only + * Tag whether is keys-only sort + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void SortBlockedToStriped( + KeyT (&keys)[ItemsPerThread], + ValueT (&values)[ItemsPerThread], + int begin_bit, + int end_bit, + ::cuda::std::bool_constant is_descending, + ::cuda::std::bool_constant is_keys_only, + DecomposerT decomposer = {}) + { + bit_ordered_type(&unsigned_keys)[ItemsPerThread] = reinterpret_cast(keys); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int KEY = 0; KEY < ItemsPerThread; KEY++) + { + unsigned_keys[KEY] = bit_ordered_conversion::to_bit_ordered(decomposer, unsigned_keys[KEY]); + } + + // Radix sorting passes + while (true) + { + int pass_bits = ::cuda::std::min(RadixBits, end_bit - begin_bit); + auto digit_extractor = + traits::template digit_extractor(begin_bit, pass_bits, decomposer); + + // Rank the blocked keys + int ranks[ItemsPerThread]; + RankKeys(unsigned_keys, ranks, digit_extractor, is_descending); + begin_bit += RadixBits; + + __syncthreads(); + + // Check if this is the last pass + if (begin_bit >= end_bit) + { + // Last pass exchanges keys through shared memory in striped arrangement + BlockExchangeKeys(temp_storage.exchange_keys).ScatterToStriped(keys, ranks); + + // Last pass exchanges through shared memory in striped arrangement + ExchangeValues(values, ranks, is_keys_only, ::cuda::std::false_type()); + + // Quit + break; + } + + // Exchange keys through shared memory in blocked arrangement + BlockExchangeKeys(temp_storage.exchange_keys).ScatterToBlocked(keys, ranks); + + // Exchange values through shared memory in blocked arrangement + ExchangeValues(values, ranks, is_keys_only, ::cuda::std::true_type()); + + __syncthreads(); + } + + // Untwiddle bits if necessary + _CCCL_PRAGMA_UNROLL_FULL() + for (int KEY = 0; KEY < ItemsPerThread; KEY++) + { + unsigned_keys[KEY] = bit_ordered_conversion::from_bit_ordered(decomposer, unsigned_keys[KEY]); + } + } + +#endif // _CCCL_DOXYGEN_INVOKED + + /// @smemstorage{BlockRadixSort} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using a private static allocation of shared memory as temporary storage. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRadixSort() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @brief Collective constructor using the specified memory allocation as temporary storage. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @param[in] temp_storage + * Reference to memory allocation having layout type TempStorage + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRadixSort(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Sorting (blocked arrangements) + //! @{ + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sort of 512 integer keys that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive keys. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer keys each + //! using BlockRadixSort = cub::BlockRadixSort; + //! + //! // Allocate shared memory for BlockRadixSort + //! __shared__ typename BlockRadixSort::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_keys[4]; + //! ... + //! + //! // Collectively sort the keys + //! BlockRadixSort(temp_storage).Sort(thread_keys); + //! + //! Suppose the set of input ``thread_keys`` across the block of threads is + //! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. + //! The corresponding output ``thread_keys`` in those threads will be + //! ``{ [0,1,2,3], [4,5,6,7], [8,9,10,11], ..., [508,509,510,511] }``. + //! @endrst + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in] begin_bit + //! **[optional]** The beginning (least-significant) bit index needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The past-the-end (most-significant) bit index needed for key comparison + _CCCL_DEVICE _CCCL_FORCEINLINE void + Sort(KeyT (&keys)[ItemsPerThread], int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + NullType values[ItemsPerThread]; + + SortBlocked(keys, values, begin_bit, end_bit, ::cuda::std::false_type(), detail::bool_constant_v); + } + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 2 keys that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 1 key. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-bits + //! :end-before: example-end keys-bits + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `(sizeof(float) + sizeof(long long int)) * 8`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + Sort(KeyT (&keys)[ItemsPerThread], DecomposerT decomposer, int begin_bit, int end_bit) + { + NullType values[ItemsPerThread]; + + SortBlocked( + keys, values, begin_bit, end_bit, ::cuda::std::false_type(), detail::bool_constant_v, decomposer); + } + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 6 keys that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 3 consecutive keys. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys + //! :end-before: example-end keys + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + Sort(KeyT (&keys)[ItemsPerThread], DecomposerT decomposer) + { + Sort(keys, decomposer, 0, detail::radix::traits_t::default_end_bit(decomposer)); + } + + //! @rst + //! Performs an ascending block-wide radix sort across a :ref:`blocked arrangement ` + //! of keys and values. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - BlockRadixSort can only accommodate one associated tile of values. To "truck along" + //! more than one tile of values, simply perform a key-value sort of the keys paired + //! with a temporary value array that enumerates the key indices. The reordered indices + //! can then be used as a gather-vector for exchanging other associated tile data through + //! shared memory. + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sort of 512 integer keys and values that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive pairs. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer keys and values each + //! using BlockRadixSort = cub::BlockRadixSort; + //! + //! // Allocate shared memory for BlockRadixSort + //! __shared__ typename BlockRadixSort::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_keys[4]; + //! int thread_values[4]; + //! ... + //! + //! // Collectively sort the keys and values among block threads + //! BlockRadixSort(temp_storage).Sort(thread_keys, thread_values); + //! + //! @endcode + //! @par + //! Suppose the set of input ``thread_keys`` across the block of threads is + //! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. The + //! corresponding output ``thread_keys`` in those threads will be + //! ``{ [0,1,2,3], [4,5,6,7], [8,9,10,11], ..., [508,509,510,511] }``. + //! + //! @endrst + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param[in] begin_bit + //! **[optional]** The beginning (least-significant) bit index needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The past-the-end (most-significant) bit index needed for key comparison + _CCCL_DEVICE _CCCL_FORCEINLINE void Sort( + KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + SortBlocked(keys, values, begin_bit, end_bit, ::cuda::std::false_type(), detail::bool_constant_v); + } + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys and values. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * BlockRadixSort can only accommodate one associated tile of values. To "truck along" + //! more than one tile of values, simply perform a key-value sort of the keys paired + //! with a temporary value array that enumerates the key indices. The reordered indices + //! can then be used as a gather-vector for exchanging other associated tile data through + //! shared memory. + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 2 keys and values that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 1 pair. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-bits + //! :end-before: example-end pairs-bits + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `(sizeof(float) + sizeof(long long int)) * 8`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + Sort( + KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], DecomposerT decomposer, int begin_bit, int end_bit) + { + SortBlocked( + keys, values, begin_bit, end_bit, ::cuda::std::false_type(), detail::bool_constant_v, decomposer); + } + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys and values. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * BlockRadixSort can only accommodate one associated tile of values. To "truck along" + //! more than one tile of values, simply perform a key-value sort of the keys paired + //! with a temporary value array that enumerates the key indices. The reordered indices + //! can then be used as a gather-vector for exchanging other associated tile data through + //! shared memory. + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 6 keys and values that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 3 consecutive pairs. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs + //! :end-before: example-end pairs + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + Sort(KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], DecomposerT decomposer) + { + Sort(keys, values, decomposer, 0, detail::radix::traits_t::default_end_bit(decomposer)); + } + + //! @rst + //! Performs a descending block-wide radix sort over a :ref:`blocked arrangement ` + //! of keys. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sort of 512 integer keys that + //! are partitioned in a [blocked arrangement](../index.html#sec5sec3) across 128 threads + //! where each thread owns 4 consecutive keys. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer keys each + //! using BlockRadixSort = cub::BlockRadixSort; + //! + //! // Allocate shared memory for BlockRadixSort + //! __shared__ typename BlockRadixSort::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_keys[4]; + //! ... + //! + //! // Collectively sort the keys + //! BlockRadixSort(temp_storage).Sort(thread_keys); + //! + //! Suppose the set of input ``thread_keys`` across the block of threads is + //! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. + //! The corresponding output ``thread_keys`` in those threads will be + //! ``{ [511,510,509,508], [11,10,9,8], [7,6,5,4], ..., [3,2,1,0] }``. + //! + //! @endrst + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in] begin_bit + //! **[optional]** The beginning (least-significant) bit index needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The past-the-end (most-significant) bit index needed for key comparison + _CCCL_DEVICE _CCCL_FORCEINLINE void + SortDescending(KeyT (&keys)[ItemsPerThread], int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + NullType values[ItemsPerThread]; + + SortBlocked(keys, values, begin_bit, end_bit, ::cuda::std::true_type(), detail::bool_constant_v); + } + + //! @rst + //! Performs a descending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 2 keys that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 1 key. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-descending-bits + //! :end-before: example-end keys-descending-bits + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `(sizeof(float) + sizeof(long long int)) * 8`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortDescending(KeyT (&keys)[ItemsPerThread], DecomposerT decomposer, int begin_bit, int end_bit) + { + NullType values[ItemsPerThread]; + + SortBlocked( + keys, values, begin_bit, end_bit, ::cuda::std::true_type(), detail::bool_constant_v, decomposer); + } + + //! @rst + //! Performs a descending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 6 keys that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 3 consecutive keys. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-descending + //! :end-before: example-end keys-descending + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortDescending(KeyT (&keys)[ItemsPerThread], DecomposerT decomposer) + { + NullType values[ItemsPerThread]; + + SortBlocked( + keys, + values, + 0, + detail::radix::traits_t::default_end_bit(decomposer), + ::cuda::std::true_type(), + detail::bool_constant_v, + decomposer); + } + + //! @rst + //! Performs a descending block-wide radix sort across a :ref:`blocked arrangement ` + //! of keys and values. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - BlockRadixSort can only accommodate one associated tile of values. To "truck along" + //! more than one tile of values, simply perform a key-value sort of the keys paired + //! with a temporary value array that enumerates the key indices. The reordered indices + //! can then be used as a gather-vector for exchanging other associated tile data through + //! shared memory. + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sort of 512 integer keys and values that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive pairs. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer keys and values each + //! using BlockRadixSort = cub::BlockRadixSort; + //! + //! // Allocate shared memory for BlockRadixSort + //! __shared__ typename BlockRadixSort::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_keys[4]; + //! int thread_values[4]; + //! ... + //! + //! // Collectively sort the keys and values among block threads + //! BlockRadixSort(temp_storage).Sort(thread_keys, thread_values); + //! + //! Suppose the set of input ``thread_keys`` across the block of threads is + //! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. The + //! corresponding output ``thread_keys`` in those threads will be + //! ``{ [511,510,509,508], [11,10,9,8], [7,6,5,4], ..., [3,2,1,0] }``. + //! + //! @endrst + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param[in] begin_bit + //! **[optional]** The beginning (least-significant) bit index needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The past-the-end (most-significant) bit index needed for key comparison + _CCCL_DEVICE _CCCL_FORCEINLINE void SortDescending( + KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + SortBlocked(keys, values, begin_bit, end_bit, ::cuda::std::true_type(), detail::bool_constant_v); + } + + //! @rst + //! Performs a descending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys and values. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * BlockRadixSort can only accommodate one associated tile of values. To "truck along" + //! more than one tile of values, simply perform a key-value sort of the keys paired + //! with a temporary value array that enumerates the key indices. The reordered indices + //! can then be used as a gather-vector for exchanging other associated tile data through + //! shared memory. + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 2 pairs that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 1 pair. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-descending-bits + //! :end-before: example-end pairs-descending-bits + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `(sizeof(float) + sizeof(long long int)) * 8`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortDescending( + KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], DecomposerT decomposer, int begin_bit, int end_bit) + { + SortBlocked( + keys, values, begin_bit, end_bit, ::cuda::std::true_type(), detail::bool_constant_v, decomposer); + } + + //! @rst + //! Performs a descending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys and values. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * BlockRadixSort can only accommodate one associated tile of values. To "truck along" + //! more than one tile of values, simply perform a key-value sort of the keys paired + //! with a temporary value array that enumerates the key indices. The reordered indices + //! can then be used as a gather-vector for exchanging other associated tile data through + //! shared memory. + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 6 keys and values that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 3 consecutive pairs. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-descending + //! :end-before: example-end pairs-descending + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortDescending(KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], DecomposerT decomposer) + { + SortBlocked( + keys, + values, + 0, + detail::radix::traits_t::default_end_bit(decomposer), + ::cuda::std::true_type(), + detail::bool_constant_v, + decomposer); + } + + //! @} + //! @name Sorting (blocked arrangement -> striped arrangement) + //! @{ + + //! @rst + //! Performs an ascending radix sort across a :ref:`blocked arrangement ` of keys, + //! leaving them in a :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sort of 512 integer keys that + //! are initially partitioned in a :ref:`blocked arrangement ` across 128 + //! threads where each thread owns 4 consecutive keys. The final partitioning is striped. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer keys each + //! using BlockRadixSort = cub::BlockRadixSort; + //! + //! // Allocate shared memory for BlockRadixSort + //! __shared__ typename BlockRadixSort::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_keys[4]; + //! ... + //! + //! // Collectively sort the keys + //! BlockRadixSort(temp_storage).SortBlockedToStriped(thread_keys); + //! + //! Suppose the set of input ``thread_keys`` across the block of threads is + //! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. + //! The corresponding output ``thread_keys`` in those threads will be + //! ``{ [0,128,256,384], [1,129,257,385], [2,130,258,386], ..., [127,255,383,511] }``. + //! + //! @endrst + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in] begin_bit + //! **[optional]** The beginning (least-significant) bit index needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The past-the-end (most-significant) bit index needed for key comparison + _CCCL_DEVICE _CCCL_FORCEINLINE void + SortBlockedToStriped(KeyT (&keys)[ItemsPerThread], int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + NullType values[ItemsPerThread]; + + SortBlockedToStriped( + keys, values, begin_bit, end_bit, ::cuda::std::false_type(), detail::bool_constant_v); + } + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys, leaving them in a + //! :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 4 keys that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 2 consecutive keys. The final partitioning is striped. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-striped-bits + //! :end-before: example-end keys-striped-bits + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `(sizeof(float) + sizeof(long long int)) * 8`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortBlockedToStriped(KeyT (&keys)[ItemsPerThread], DecomposerT decomposer, int begin_bit, int end_bit) + { + NullType values[ItemsPerThread]; + + SortBlockedToStriped( + keys, values, begin_bit, end_bit, ::cuda::std::false_type(), detail::bool_constant_v, decomposer); + } + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys, leaving them in a + //! :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 6 keys that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 3 consecutive keys. The final partitioning is striped. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-striped + //! :end-before: example-end keys-striped + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortBlockedToStriped(KeyT (&keys)[ItemsPerThread], DecomposerT decomposer) + { + NullType values[ItemsPerThread]; + + SortBlockedToStriped( + keys, + values, + 0, + detail::radix::traits_t::default_end_bit(decomposer), + ::cuda::std::false_type(), + detail::bool_constant_v, + decomposer); + } + + //! @rst + //! Performs an ascending radix sort across a :ref:`blocked arrangement ` of keys and + //! values, leaving them in a :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - BlockRadixSort can only accommodate one associated tile of values. To "truck along" + //! more than one tile of values, simply perform a key-value sort of the keys paired + //! with a temporary value array that enumerates the key indices. The reordered indices + //! can then be used as a gather-vector for exchanging other associated tile data through + //! shared memory. + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sort of 512 integer keys and values that + //! are initially partitioned in a [blocked arrangement](../index.html#sec5sec3) across 128 + //! threads where each thread owns 4 consecutive pairs. The final partitioning is striped. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer keys and values each + //! using BlockRadixSort = cub::BlockRadixSort; + //! + //! // Allocate shared memory for BlockRadixSort + //! __shared__ typename BlockRadixSort::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_keys[4]; + //! int thread_values[4]; + //! ... + //! + //! // Collectively sort the keys and values among block threads + //! BlockRadixSort(temp_storage).SortBlockedToStriped(thread_keys, thread_values); + //! + //! Suppose the set of input ``thread_keys`` across the block of threads is + //! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. + //! The corresponding output ``thread_keys`` in those threads will be + //! ``{ [0,128,256,384], [1,129,257,385], [2,130,258,386], ..., [127,255,383,511] }``. + //! + //! @endrst + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param[in] begin_bit + //! **[optional]** The beginning (least-significant) bit index needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The past-the-end (most-significant) bit index needed for key comparison + _CCCL_DEVICE _CCCL_FORCEINLINE void SortBlockedToStriped( + KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + SortBlockedToStriped( + keys, values, begin_bit, end_bit, ::cuda::std::false_type(), detail::bool_constant_v); + } + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys and values, leaving them in a + //! :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 4 pairs that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 2 consecutive pairs. The final partitioning is striped. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-striped-bits + //! :end-before: example-end pairs-striped-bits + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `(sizeof(float) + sizeof(long long int)) * 8`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortBlockedToStriped( + KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], DecomposerT decomposer, int begin_bit, int end_bit) + { + SortBlockedToStriped( + keys, values, begin_bit, end_bit, ::cuda::std::false_type(), detail::bool_constant_v, decomposer); + } + + //! @rst + //! Performs an ascending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys and values, leaving them in a + //! :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 6 pairs that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 3 consecutive pairs. The final partitioning is striped. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-striped + //! :end-before: example-end pairs-striped + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortBlockedToStriped(KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], DecomposerT decomposer) + { + SortBlockedToStriped( + keys, + values, + 0, + detail::radix::traits_t::default_end_bit(decomposer), + ::cuda::std::false_type(), + detail::bool_constant_v, + decomposer); + } + + //! @rst + //! Performs a descending radix sort across a :ref:`blocked arrangement ` + //! of keys, leaving them in a :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sort of 512 integer keys that + //! are initially partitioned in a :ref:`blocked arrangement ` across 128 + //! threads where each thread owns 4 consecutive keys. The final partitioning is striped. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer keys each + //! using BlockRadixSort = cub::BlockRadixSort; + //! + //! // Allocate shared memory for BlockRadixSort + //! __shared__ typename BlockRadixSort::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_keys[4]; + //! ... + //! + //! // Collectively sort the keys + //! BlockRadixSort(temp_storage).SortBlockedToStriped(thread_keys); + //! + //! Suppose the set of input ``thread_keys`` across the block of threads is + //! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. + //! The corresponding output ``thread_keys`` in those threads will be + //! ``{ [511,383,255,127], [386,258,130,2], [385,257,128,1], ..., [384,256,128,0] }``. + //! + //! @endrst + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in] begin_bit + //! **[optional]** The beginning (least-significant) bit index needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The past-the-end (most-significant) bit index needed for key comparison + _CCCL_DEVICE _CCCL_FORCEINLINE void + SortDescendingBlockedToStriped(KeyT (&keys)[ItemsPerThread], int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + NullType values[ItemsPerThread]; + + SortBlockedToStriped(keys, values, begin_bit, end_bit, ::cuda::std::true_type(), detail::bool_constant_v); + } + + //! @rst + //! Performs a descending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys, leaving them in a + //! :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 4 keys that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 2 consecutive keys. The final partitioning is striped. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-striped-descending-bits + //! :end-before: example-end keys-striped-descending-bits + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `(sizeof(float) + sizeof(long long int)) * 8`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortDescendingBlockedToStriped(KeyT (&keys)[ItemsPerThread], DecomposerT decomposer, int begin_bit, int end_bit) + { + NullType values[ItemsPerThread]; + + SortBlockedToStriped( + keys, values, begin_bit, end_bit, ::cuda::std::true_type(), detail::bool_constant_v, decomposer); + } + + //! @rst + //! Performs a descending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys, leaving them in a + //! :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 6 keys that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 3 consecutive keys. The final partitioning is striped. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-striped-descending + //! :end-before: example-end keys-striped-descending + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortDescendingBlockedToStriped(KeyT (&keys)[ItemsPerThread], DecomposerT decomposer) + { + NullType values[ItemsPerThread]; + + SortBlockedToStriped( + keys, + values, + 0, + detail::radix::traits_t::default_end_bit(decomposer), + ::cuda::std::true_type(), + detail::bool_constant_v, + decomposer); + } + + //! @rst + //! Performs a descending radix sort across a :ref:`blocked arrangement ` + //! of keys and values, leaving them in a :ref:`striped arrangement ` + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - BlockRadixSort can only accommodate one associated tile of values. To "truck along" + //! more than one tile of values, simply perform a key-value sort of the keys paired + //! with a temporary value array that enumerates the key indices. The reordered indices + //! can then be used as a gather-vector for exchanging other associated tile data through + //! shared memory. + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sort of 512 integer keys and values that + //! are initially partitioned in a :ref:`blocked arrangement ` across 128 + //! threads where each thread owns 4 consecutive pairs. The final partitioning is striped. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockRadixSort for a 1D block of 128 threads owning 4 integer keys and values each + //! using BlockRadixSort = cub::BlockRadixSort; + //! + //! // Allocate shared memory for BlockRadixSort + //! __shared__ typename BlockRadixSort::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_keys[4]; + //! int thread_values[4]; + //! ... + //! + //! // Collectively sort the keys and values among block threads + //! BlockRadixSort(temp_storage).SortBlockedToStriped(thread_keys, thread_values); + //! + //! Suppose the set of input ``thread_keys`` across the block of threads is + //! ``{ [0,511,1,510], [2,509,3,508], [4,507,5,506], ..., [254,257,255,256] }``. + //! The corresponding output ``thread_keys`` in those threads will be + //! ``{ [511,383,255,127], [386,258,130,2], [385,257,128,1], ..., [384,256,128,0] }``. + //! + //! @endrst + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param[in] begin_bit + //! **[optional]** The beginning (least-significant) bit index needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The past-the-end (most-significant) bit index needed for key comparison + _CCCL_DEVICE _CCCL_FORCEINLINE void SortDescendingBlockedToStriped( + KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + SortBlockedToStriped(keys, values, begin_bit, end_bit, ::cuda::std::true_type(), detail::bool_constant_v); + } + + //! @rst + //! Performs a descending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys and values, leaving them in a + //! :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 4 keys and values that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 2 consecutive pairs. The final partitioning is striped. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-striped-descending-bits + //! :end-before: example-end pairs-striped-descending-bits + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `(sizeof(float) + sizeof(long long int)) * 8`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortDescendingBlockedToStriped( + KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], DecomposerT decomposer, int begin_bit, int end_bit) + { + SortBlockedToStriped( + keys, values, begin_bit, end_bit, ::cuda::std::true_type(), detail::bool_constant_v, decomposer); + } + + //! @rst + //! Performs a descending block-wide radix sort over a + //! :ref:`blocked arrangement ` of keys and values, leaving them in a + //! :ref:`striped arrangement `. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * @granularity + //! * @smemreuse + //! + //! Snippet + //! ========================================================================== + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The code snippet below illustrates a sort of 6 keys and values that + //! are partitioned in a :ref:`blocked arrangement ` across 2 threads + //! where each thread owns 3 consecutive pairs. The final partitioning is striped. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-striped-descending + //! :end-before: example-end pairs-striped-descending + //! + //! @endrst + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in,out] keys + //! Keys to sort + //! + //! @param[in,out] values + //! Values to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + template + _CCCL_DEVICE _CCCL_FORCEINLINE // + ::cuda::std::enable_if_t< // + !::cuda::std::is_convertible_v> + SortDescendingBlockedToStriped(KeyT (&keys)[ItemsPerThread], ValueT (&values)[ItemsPerThread], DecomposerT decomposer) + { + SortBlockedToStriped( + keys, + values, + 0, + detail::radix::traits_t::default_end_bit(decomposer), + ::cuda::std::true_type(), + detail::bool_constant_v, + decomposer); + } + + //@} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_raking_layout.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_raking_layout.cuh new file mode 100644 index 00000000..0cf6fc4a --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_raking_layout.cuh @@ -0,0 +1,124 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::BlockRakingLayout provides a conflict-free shared memory layout abstraction for warp-raking + * across thread block data. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! BlockRakingLayout provides a conflict-free shared memory layout abstraction for 1D raking across thread block data. +//! +//! Overview +//! ++++++++++++++++++++++++++ +//! +//! This type facilitates a shared memory usage pattern where a block of CUDA +//! threads places elements into shared memory and then reduces the active +//! parallelism to one "raking" warp of threads for serially aggregating consecutive +//! sequences of shared items. Padding is inserted to eliminate bank conflicts +//! (for most data types). +//! +//! @endrst +//! +//! @tparam T +//! The data type to be exchanged. +//! +//! @tparam ThreadsPerBlock +//! The thread block size in threads. +//! +template +struct BlockRakingLayout +{ + //--------------------------------------------------------------------- + // Constants and type definitions + //--------------------------------------------------------------------- + + /// The total number of elements that need to be cooperatively reduced + static constexpr int SHARED_ELEMENTS = ThreadsPerBlock; + + /// Maximum number of warp-synchronous raking threads + static constexpr int MAX_RAKING_THREADS = ::cuda::std::min(ThreadsPerBlock, detail::warp_threads); + + /// Number of raking elements per warp-synchronous raking thread (rounded up) + static constexpr int SEGMENT_LENGTH = (SHARED_ELEMENTS + MAX_RAKING_THREADS - 1) / MAX_RAKING_THREADS; + + /// Never use a raking thread that will have no valid data (e.g., when ThreadsPerBlock is 62 and SEGMENT_LENGTH is 2, + /// we should only use 31 raking threads) + static constexpr int RAKING_THREADS = (SHARED_ELEMENTS + SEGMENT_LENGTH - 1) / SEGMENT_LENGTH; + + /// Whether we will have bank conflicts (technically we should find out if the GCD is > 1) + static constexpr bool HAS_CONFLICTS = (detail::smem_banks % SEGMENT_LENGTH == 0); + + /// Degree of bank conflicts (e.g., 4-way) + static constexpr int CONFLICT_DEGREE = + (HAS_CONFLICTS) ? (MAX_RAKING_THREADS * SEGMENT_LENGTH) / detail::smem_banks : 1; + + /// Pad each segment length with one element if segment length is not relatively prime to warp size and can't be + /// optimized as a vector load + static constexpr bool USE_SEGMENT_PADDING = ((SEGMENT_LENGTH & 1) == 0) && (SEGMENT_LENGTH > 2); + + /// Total number of elements in the raking grid + static constexpr int GRID_ELEMENTS = RAKING_THREADS * (SEGMENT_LENGTH + USE_SEGMENT_PADDING); + + /// Whether or not we need bounds checking during raking (the number of reduction elements is not a multiple of the + /// number of raking threads) + static constexpr int UNGUARDED = (SHARED_ELEMENTS % RAKING_THREADS == 0); + + /** + * @brief Shared memory storage type + */ + struct __align__(16) _TempStorage + { + T buff[BlockRakingLayout::GRID_ELEMENTS]; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + /** + * @brief Returns the location for the calling thread to place data into the grid + */ + static _CCCL_DEVICE _CCCL_FORCEINLINE T* PlacementPtr(TempStorage& temp_storage, unsigned int linear_tid) + { + // Offset for partial + unsigned int offset = linear_tid; + + // Add in one padding element for every segment + if (USE_SEGMENT_PADDING > 0) + { + offset += offset / SEGMENT_LENGTH; + } + + // Incorporating a block of padding partials every shared memory segment + return temp_storage.Alias().buff + offset; + } + + /** + * @brief Returns the location for the calling thread to begin sequential raking + */ + static _CCCL_DEVICE _CCCL_FORCEINLINE T* RakingPtr(TempStorage& temp_storage, unsigned int linear_tid) + { + return temp_storage.Alias().buff + (linear_tid * (SEGMENT_LENGTH + USE_SEGMENT_PADDING)); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_reduce.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_reduce.cuh new file mode 100644 index 00000000..992d213b --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_reduce.cuh @@ -0,0 +1,711 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! The cub::BlockReduce class provides :ref:`collective ` methods for +//! computing a parallel reduction of items partitioned across a CUDA thread block. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Algorithmic variants + ******************************************************************************/ + +//! BlockReduceAlgorithm enumerates alternative algorithms for parallel reduction across a CUDA thread +//! block. +enum BlockReduceAlgorithm +{ + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! An efficient "raking" reduction algorithm that only supports commutative reduction operators + //! (true for most operations, e.g., addition). + //! + //! Execution is comprised of three phases: + //! #. Upsweep sequential reduction in registers (if threads contribute more than one input each). + //! Threads in warps other than the first warp place their partial reductions into shared + //! memory. + //! #. Upsweep sequential reduction in shared memory. Threads within the first warp continue to + //! accumulate by raking across segments of shared partial reductions + //! #. A warp-synchronous Kogge-Stone style reduction within the raking warp. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - This variant performs less communication than BLOCK_REDUCE_RAKING_NON_COMMUTATIVE and is + //! preferable when the reduction operator is commutative. This variant applies fewer reduction + //! operators than BLOCK_REDUCE_WARP_REDUCTIONS, and can provide higher overall throughput across + //! the GPU when suitably occupied. However, turn-around latency may be higher than to + //! BLOCK_REDUCE_WARP_REDUCTIONS and thus less-desirable when the GPU is under-occupied. + //! + //! @endrst + BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! An efficient "raking" reduction algorithm that supports commutative (e.g., addition) and + //! non-commutative (e.g., string concatenation) reduction operators. @blocked. + //! + //! Execution is comprised of three phases: + //! #. Upsweep sequential reduction in registers (if threads contribute more than one input each). + //! Each thread then places the partial reduction of its item(s) into shared memory. + //! #. Upsweep sequential reduction in shared memory. Threads within a single warp rake across + //! segments of shared partial reductions. + //! #. A warp-synchronous Kogge-Stone style reduction within the raking warp. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - This variant performs more communication than BLOCK_REDUCE_RAKING and is only preferable when + //! the reduction operator is non-commutative. This variant applies fewer reduction operators than + //! BLOCK_REDUCE_WARP_REDUCTIONS, and can provide higher overall throughput across the GPU when + //! suitably occupied. However, turn-around latency may be higher than to + //! BLOCK_REDUCE_WARP_REDUCTIONS and thus less-desirable when the GPU is under-occupied. + //! + //! @endrst + BLOCK_REDUCE_RAKING, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A quick "tiled warp-reductions" reduction algorithm that supports commutative (e.g., addition) + //! and non-commutative (e.g., string concatenation) reduction operators. + //! + //! Execution is comprised of four phases: + //! #. Upsweep sequential reduction in registers (if threads contribute more than one input each). + //! Each thread then places the partial reduction of its item(s) into shared memory. + //! #. Compute a shallow, but inefficient warp-synchronous Kogge-Stone style reduction within + //! each warp. + //! #. A propagation phase where the warp reduction outputs in each warp are updated with the + //! aggregate from each preceding warp. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - This variant applies more reduction operators than BLOCK_REDUCE_RAKING or + //! BLOCK_REDUCE_RAKING_NON_COMMUTATIVE, which may result in lower overall throughput across the + //! GPU. However turn-around latency may be lower and thus useful when the GPU is under-occupied. + //! + //! @endrst + BLOCK_REDUCE_WARP_REDUCTIONS, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A quick "tiled warp-reductions" reduction algorithm that supports commutative (e.g., addition) + //! and non-commutative (e.g., string concatenation) reduction operators. This variant uses atomic + //! operations to reduce the warp-wide reduction results, making it non-deterministic, i.e. the + //! order of reduction operations is not guaranteed to be the same across different invocations of + //! the same kernel. + //! + //! Execution is comprised of three phases: + //! #. Upsweep sequential reduction in registers (if threads contribute more than one input each). + //! Each thread then places the partial reduction of its item(s) into shared memory. + //! #. Compute a shallow, but non work-efficient warp-synchronous Kogge-Stone style reduction + //! within each warp. + //! #. Lane 0 of warp 0 stores its warp aggregate, while lane 0 of other warps use atomic + //! operations to accumulate their warp aggregates into a shared location, making the final + //! order non-deterministic. + //! #. The final block-wide result is available to all threads. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - This variant applies more reduction operators than BLOCK_REDUCE_RAKING or + //! BLOCK_REDUCE_RAKING_NON_COMMUTATIVE, which may result in lower overall throughput across the + //! GPU. However turn-around latency may be lower and thus useful when the GPU is under-occupied. + //! + //! @endrst + BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC, +}; + +#if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED) +namespace detail +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(BlockReduceAlgorithm algo) noexcept +{ + switch (algo) + { + case BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY: + return "BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY"; + case BLOCK_REDUCE_RAKING: + return "BLOCK_REDUCE_RAKING"; + case BLOCK_REDUCE_WARP_REDUCTIONS: + return "BLOCK_REDUCE_WARP_REDUCTIONS"; + case BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC: + return "BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC"; + } + return ""; +} +} // namespace detail + +inline ::std::ostream& operator<<(::std::ostream& os, BlockReduceAlgorithm algo) +{ + return os << CUB_NS_QUALIFIER::detail::to_string(algo); +} +#endif // _CCCL_HOSTED() && !_CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END + +#if __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) +template <::cuda::std::same_as CharT> +struct std::formatter : formatter +{ + template + auto format(const CUB_NS_QUALIFIER::BlockReduceAlgorithm& algo, FmtCtx& ctx) const + { + return formatter::format(CUB_NS_QUALIFIER::detail::to_string(algo), ctx); + } +}; +#endif // __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) + +CUB_NAMESPACE_BEGIN + +//! @rst +//! The BlockReduce class provides :ref:`collective ` methods for computing a +//! parallel reduction of items partitioned across a CUDA thread block. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - A `reduction `_ (or *fold*) uses a +//! binary combining operator to compute a single aggregate from a list of input elements. +//! - @rowmajor +//! - BlockReduce can be optionally specialized by algorithm to accommodate different +//! latency/throughput workload profiles: +//! +//! #. :cpp:enumerator:`cub::BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY`: +//! An efficient "raking" reduction algorithm that only supports commutative reduction operators. +//! #. :cpp:enumerator:`cub::BLOCK_REDUCE_RAKING`: +//! An efficient "raking" reduction algorithm that supports commutative and non-commutative +//! reduction operators. +//! #. :cpp:enumerator:`cub::BLOCK_REDUCE_WARP_REDUCTIONS`: +//! A quick "tiled warp-reductions" reduction algorithm that supports commutative and +//! non-commutative reduction operators. +//! #. :cpp:enumerator:`cub::BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC`: +//! A quick "tiled warp-reductions" reduction algorithm that supports commutative and +//! non-commutative reduction operators. This variant uses atomic operations to reduce the +//! warp-wide reduction results, making it non-deterministic, i.e. the order of reduction +//! operations is not guaranteed to be the same across different invocations of the same +//! kernel. +//! +//! Performance Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - @granularity +//! - Very efficient (only one synchronization barrier). +//! - Incurs zero bank conflicts for most types +//! - Computation is slightly more efficient (i.e., having lower instruction overhead) for: +//! - Summation (vs. generic reduction) +//! - ``BLOCK_THREADS`` is a multiple of the architecture's warp size +//! - Every thread has a valid input (i.e., full vs. partial-tiles) +//! - See cub::BlockReduceAlgorithm for performance details regarding algorithmic alternatives +//! +//! A Simple Example +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @blockcollective{BlockReduce} +//! +//! The code snippet below illustrates a sum reduction of 512 integer items that are partitioned in +//! a :ref:`blocked arrangement ` across 128 threads where each thread +//! owns 4 consecutive items. +//! +//! .. code-block:: c++ +//! +//! #include // or equivalently +//! +//! __global__ void ExampleKernel(...) +//! { +//! // Specialize BlockReduce for a 1D block of 128 threads of type int +//! using BlockReduce = cub::BlockReduce; +//! +//! // Allocate shared memory for BlockReduce +//! __shared__ typename BlockReduce::TempStorage temp_storage; +//! +//! // Obtain a segment of consecutive items that are blocked across threads +//! int thread_data[4]; +//! ... +//! +//! // Compute the block-wide sum for thread0 +//! int aggregate = BlockReduce(temp_storage).Sum(thread_data); +//! } +//! +//! Re-using dynamically allocating shared memory +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of dynamically shared +//! memory with BlockReduce and how to re-purpose the same memory region. +//! +//! @endrst +//! +//! @tparam T +//! Data type being reduced +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam Algorithm +//! **[optional]** cub::BlockReduceAlgorithm enumerator specifying the underlying algorithm to use +//! (default: cub::BLOCK_REDUCE_WARP_REDUCTIONS) +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockReduce +{ +private: + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + using WarpReductions = detail::BlockReduceWarpReductions; + using WarpReductionsNondeterministic = detail::BlockReduceWarpReductions; + using RakingCommutativeOnly = detail::BlockReduceRakingCommutativeOnly; + using Raking = detail::BlockReduceRaking; + + /// Internal specialization type + using InternalBlockReduce = + ::cuda::std::_If>>; // BlockReduceRaking + + /// Shared memory storage layout type for BlockReduce + using _TempStorage = typename InternalBlockReduce::TempStorage; + + /// Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + unsigned int linear_tid; + +public: + /// @smemstorage{BlockReduce} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using a private static allocation of shared memory as temporary + //! storage. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + _CCCL_DEVICE _CCCL_FORCEINLINE BlockReduce() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @brief Collective constructor using the specified memory allocation as temporary storage. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @param[in] temp_storage + * Reference to memory allocation having layout type TempStorage + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockReduce(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Generic reductions + //! @{ + + //! @rst + //! Computes a block-wide reduction for thread\ :sub:`0` using the specified binary reduction functor. + //! Each thread contributes one input element. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The return value is undefined in threads other than thread\ :sub:`0`. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a max reduction of 128 integer items that are partitioned + //! across 128 threads. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockReduce for a 1D block of 128 threads of type int + //! using BlockReduce = cub::BlockReduce; + //! + //! // Allocate shared memory for BlockReduce + //! __shared__ typename BlockReduce::TempStorage temp_storage; + //! + //! // Each thread obtains an input item + //! int thread_data; + //! ... + //! + //! // Compute the block-wide max for thread0 + //! int aggregate = BlockReduce(temp_storage).Reduce(thread_data, cuda::maximum<>{}); + //! } + //! + //! @endrst + //! + //! @tparam ReductionOp + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input + //! + //! @param[in] reduction_op + //! Binary reduction functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Reduce(T input, ReductionOp reduction_op) + { + return InternalBlockReduce(temp_storage).template Reduce(input, BLOCK_THREADS, reduction_op); + } + + //! @rst + //! Computes a block-wide reduction for thread\ :sub:`0` using the specified binary reduction + //! functor. Each thread contributes an array of consecutive input elements. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The return value is undefined in threads other than thread\ :sub:`0`. + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a max reduction of 512 integer items that are partitioned in a + //! :ref:`blocked arrangement ` across 128 threads where each thread owns + //! 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockReduce for a 1D block of 128 threads of type int + //! using BlockReduce = cub::BlockReduce; + //! + //! // Allocate shared memory for BlockReduce + //! __shared__ typename BlockReduce::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Compute the block-wide max for thread0 + //! int aggregate = BlockReduce(temp_storage).Reduce(thread_data, cuda::maximum<>{}); + //! } + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ReductionOp + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] inputs + //! Calling thread's input segment + //! + //! @param[in] reduction_op + //! Binary reduction functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Reduce(T (&inputs)[ITEMS_PER_THREAD], ReductionOp reduction_op) + { + // Reduce partials + T partial = cub::ThreadReduce(inputs, reduction_op); + return Reduce(partial, reduction_op); + } + + //! @rst + //! Computes a block-wide reduction for thread\ :sub:`0` using the specified binary reduction + //! functor. The first ``num_valid`` threads each contribute one input element. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The return value is undefined in threads other than thread0. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a max reduction of a partially-full tile of integer items + //! that are partitioned across 128 threads. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int num_valid, ...) + //! { + //! // Specialize BlockReduce for a 1D block of 128 threads of type int + //! using BlockReduce = cub::BlockReduce; + //! + //! // Allocate shared memory for BlockReduce + //! __shared__ typename BlockReduce::TempStorage temp_storage; + //! + //! // Each thread obtains an input item + //! int thread_data; + //! if (threadIdx.x < num_valid) thread_data = ... + //! + //! // Compute the block-wide max for thread0 + //! int aggregate = BlockReduce(temp_storage).Reduce(thread_data, cuda::maximum<>{}, num_valid); + //! } + //! + //! @endrst + //! + //! @tparam ReductionOp + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] num_valid + //! Number of threads containing valid elements (may be less than BLOCK_THREADS) + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Reduce(T input, ReductionOp reduction_op, int num_valid) + { + // Determine if we skip bounds checking + if (num_valid >= BLOCK_THREADS) + { + return InternalBlockReduce(temp_storage).template Reduce(input, num_valid, reduction_op); + } + else + { + return InternalBlockReduce(temp_storage).template Reduce(input, num_valid, reduction_op); + } + } + + //! @} + //! @name Summation reductions + //! @{ + + //! @rst + //! Computes a block-wide reduction for thread\ :sub:`0` using addition (+) as the reduction operator. + //! Each thread contributes one input element. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The return value is undefined in threads other than thread\ :sub:`0`. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sum reduction of 128 integer items that are partitioned + //! across 128 threads. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockReduce for a 1D block of 128 threads of type int + //! using BlockReduce = cub::BlockReduce; + //! + //! // Allocate shared memory for BlockReduce + //! __shared__ typename BlockReduce::TempStorage temp_storage; + //! + //! // Each thread obtains an input item + //! int thread_data; + //! ... + //! + //! // Compute the block-wide sum for thread0 + //! int aggregate = BlockReduce(temp_storage).Sum(thread_data); + //! } + //! + //! @endrst + //! + //! @param[in] input + //! Calling thread's input + _CCCL_DEVICE _CCCL_FORCEINLINE T Sum(T input) + { + return InternalBlockReduce(temp_storage).template Sum(input, BLOCK_THREADS); + } + + //! @rst + //! Computes a block-wide reduction for thread0 using addition (+) as the reduction + //! operator. Each thread contributes an array of consecutive input elements. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The return value is undefined in threads other than thread\ :sub:`0`. + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sum reduction of 512 integer items that are partitioned in a + //! :ref:`blocked arrangement ` across 128 threads where each thread owns + //! 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockReduce for a 1D block of 128 threads of type int + //! using BlockReduce = cub::BlockReduce; + //! + //! // Allocate shared memory for BlockReduce + //! __shared__ typename BlockReduce::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Compute the block-wide sum for thread0 + //! int aggregate = BlockReduce(temp_storage).Sum(thread_data); + //! } + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @param[in] inputs + //! Calling thread's input segment + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Sum(T (&inputs)[ITEMS_PER_THREAD]) + { + // Reduce partials + T partial = cub::ThreadReduce(inputs, ::cuda::std::plus<>{}); + return Sum(partial); + } + + //! @rst + //! Computes a block-wide reduction for thread\ :sub:`0` using addition (+) as the reduction + //! operator. The first ``num_valid`` threads each contribute one input element. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The return value is undefined in threads other than thread\ :sub:`0`. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a sum reduction of a partially-full tile of integer items + //! that are partitioned across 128 threads. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int num_valid, ...) + //! { + //! // Specialize BlockReduce for a 1D block of 128 threads of type int + //! using BlockReduce = cub::BlockReduce; + //! + //! // Allocate shared memory for BlockReduce + //! __shared__ typename BlockReduce::TempStorage temp_storage; + //! + //! // Each thread obtains an input item (up to num_items) + //! int thread_data; + //! if (threadIdx.x < num_valid) + //! thread_data = ... + //! + //! // Compute the block-wide sum for thread0 + //! int aggregate = BlockReduce(temp_storage).Sum(thread_data, num_valid); + //! } + //! + //! @endrst + //! + //! @param[in] input + //! Calling thread's input + //! + //! @param[in] num_valid + //! Number of threads containing valid elements (may be less than BLOCK_THREADS) + _CCCL_DEVICE _CCCL_FORCEINLINE T Sum(T input, int num_valid) + { + // Determine if we skip bounds checking + if (num_valid >= BLOCK_THREADS) + { + return InternalBlockReduce(temp_storage).template Sum(input, num_valid); + } + else + { + return InternalBlockReduce(temp_storage).template Sum(input, num_valid); + } + } + + //! @} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_run_length_decode.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_run_length_decode.cuh new file mode 100644 index 00000000..1ce6cbcb --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_run_length_decode.cuh @@ -0,0 +1,436 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! The BlockRunLengthDecode class supports decoding a run-length encoded array of items. That +//! is, given the two arrays ``run_value[N]`` and ``run_lengths[N]``, ``run_value[i]`` is repeated ``run_lengths[i]`` +//! many times in the output array. Due to the nature of the run-length decoding algorithm +//! ("decompression"), the output size of the run-length decoded array is runtime-dependent and +//! potentially without any upper bound. To address this, BlockRunLengthDecode allows retrieving a +//! "window" from the run-length decoded array. The window's offset can be specified and +//! BLOCK_THREADS * DecodedItemsPerThread (i.e., referred to as window_size) decoded items from +//! the specified window will be returned. +//! +//! .. note:: +//! +//! Trailing runs of length 0 are supported (i.e., they may only appear at the end of the run_lengths array). +//! A run of length zero may not be followed by a run length that is not zero. +//! +//! .. code-block:: c++ +//! +//! __global__ void ExampleKernel(...) +//! { +//! // Specialising BlockRunLengthDecode to run-length decode items of type uint64_t +//! using RunItemT = uint64_t; +//! // Type large enough to index into the run-length decoded array +//! using RunLengthT = uint32_t; +//! +//! // Specialising BlockRunLengthDecode for a 1D block of 128 threads +//! constexpr int BlockDimX = 128; +//! // Specialising BlockRunLengthDecode to have each thread contribute 2 run-length encoded runs +//! constexpr int RunsPerThread = 2; +//! // Specialising BlockRunLengthDecode to have each thread hold 4 run-length decoded items +//! constexpr int DecodedItemsPerThread = 4; +//! +//! // Specialize BlockRunLengthDecode for a 1D block of 128 threads owning 4 integer items each +//! using BlockRunLengthDecodeT = +//! cub::BlockRunLengthDecode; +//! +//! // Allocate shared memory for BlockRunLengthDecode +//! __shared__ typename BlockRunLengthDecodeT::TempStorage temp_storage; +//! +//! // The run-length encoded items and how often they shall be repeated in the run-length decoded output +//! RunItemT run_values[RunsPerThread]; +//! RunLengthT run_lengths[RunsPerThread]; +//! ... +//! +//! // Initialize the BlockRunLengthDecode with the runs that we want to run-length decode +//! uint32_t total_decoded_size = 0; +//! BlockRunLengthDecodeT block_rld(temp_storage, run_values, run_lengths, total_decoded_size); +//! +//! // Run-length decode ("decompress") the runs into a window buffer of limited size. This is repeated until all +//! runs +//! // have been decoded. +//! uint32_t decoded_window_offset = 0U; +//! while (decoded_window_offset < total_decoded_size) +//! { +//! RunLengthT relative_offsets[DecodedItemsPerThread]; +//! RunItemT decoded_items[DecodedItemsPerThread]; +//! +//! // The number of decoded items that are valid within this window (aka pass) of run-length decoding +//! uint32_t num_valid_items = total_decoded_size - decoded_window_offset; +//! block_rld.RunLengthDecode(decoded_items, relative_offsets, decoded_window_offset); +//! +//! decoded_window_offset += BlockDimX * DecodedItemsPerThread; +//! +//! ... +//! } +//! } +//! +//! Suppose the set of input ``run_values`` across the block of threads is +//! ``{ [0, 1], [2, 3], [4, 5], [6, 7], ..., [254, 255] }`` and +//! ``run_lengths`` is ``{ [1, 2], [3, 4], [5, 1], [2, 3], ..., [5, 1] }``. +//! The corresponding output ``decoded_items`` in those threads will be +//! ``{ [0, 1, 1, 2], [2, 2, 3, 3], [3, 3, 4, 4], [4, 4, 4, 5], ..., [169, 169, 170, 171] }`` +//! and ``relative_offsets`` will be +//! ``{ [0, 0, 1, 0], [1, 2, 0, 1], [2, 3, 0, 1], [2, 3, 4, 0], ..., [3, 4, 0, 0] }`` during the +//! first iteration of the while loop. +//! +//! @endrst +//! +//! @tparam ItemT +//! The data type of the items being run-length decoded +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam RunsPerThread +//! The number of consecutive runs that each thread contributes +//! +//! @tparam DecodedItemsPerThread +//! The maximum number of decoded items that each thread holds +//! +//! @tparam DecodedOffsetT +//! Type used to index into the block's decoded items (large enough to hold the sum over all the +//! runs' lengths) +//! +//! @tparam BlockDimY +//! The thread block length in threads along the Y dimension +//! +//! @tparam BlockDimZ +//! The thread block length in threads along the Z dimension +template +class BlockRunLengthDecode +{ + //--------------------------------------------------------------------- + // CONFIGS & TYPE ALIASES + //--------------------------------------------------------------------- + +private: + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + /// The number of runs that the block decodes (out-of-bounds items may be padded with run lengths of '0') + static constexpr int BLOCK_RUNS = BLOCK_THREADS * RunsPerThread; + + /// BlockScan used to determine the beginning of each run (i.e., prefix sum over the runs' length) + using RunOffsetScanT = BlockScan; + + /// Type used to index into the block's runs + using RunOffsetT = uint32_t; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + /// Shared memory type required by this thread block + union _TempStorage + { + typename RunOffsetScanT::TempStorage offset_scan; + struct + { + ItemT run_values[BLOCK_RUNS]; + DecodedOffsetT run_offsets[BLOCK_RUNS]; + } runs; + }; // union TempStorage +#endif // _CCCL_DOXYGEN_INVOKED + + /// Internal storage allocator (used when the user does not provide pre-allocated shared memory) + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + uint32_t linear_tid; + +public: + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // CONSTRUCTOR + //--------------------------------------------------------------------- + + //! @brief Constructor specialised for user-provided temporary storage, initializing using the runs' lengths. + //! The algorithm's temporary storage may not be repurposed between the constructor call and subsequent + //! `RunLengthDecode` calls. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRunLengthDecode( + TempStorage& temp_storage, + ItemT (&run_values)[RunsPerThread], + RunLengthT (&run_lengths)[RunsPerThread], + TotalDecodedSizeT& total_decoded_size) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + { + InitWithRunLengths(run_values, run_lengths, total_decoded_size); + } + + //! @brief Constructor specialised for user-provided temporary storage, initializing using the runs' offsets. + //! The algorithm's temporary storage may not be repurposed between the constructor call and subsequent + //! `RunLengthDecode` calls. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRunLengthDecode( + TempStorage& temp_storage, ItemT (&run_values)[RunsPerThread], UserRunOffsetT (&run_offsets)[RunsPerThread]) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + { + InitWithRunOffsets(run_values, run_offsets); + } + + /** + * \brief Constructor specialised for static temporary storage, initializing using the runs' lengths. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE BlockRunLengthDecode( + ItemT (&run_values)[RunsPerThread], RunLengthT (&run_lengths)[RunsPerThread], TotalDecodedSizeT& total_decoded_size) + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + { + InitWithRunLengths(run_values, run_lengths, total_decoded_size); + } + + /** + * \brief Constructor specialised for static temporary storage, initializing using the runs' offsets. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE + BlockRunLengthDecode(ItemT (&run_values)[RunsPerThread], UserRunOffsetT (&run_offsets)[RunsPerThread]) + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + { + InitWithRunOffsets(run_values, run_offsets); + } + +private: + /** + * @brief Returns the offset of the first value within @p input which compares greater than + * @p val. This version takes @p MAX_NUM_ITEMS, an upper bound of the array size, which will + * be used to determine the number of binary search iterations at compile time. + * + * @param[in] input + * Input sequence + * + * @param[in] num_items + * Input sequence length + * + * @param[in] val + * Search key + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE OffsetT StaticUpperBound(InputIteratorT input, OffsetT num_items, T val) + { + OffsetT lower_bound = 0; + OffsetT upper_bound = num_items; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i <= Log2::VALUE; i++) + { + OffsetT mid = cub::MidPoint(lower_bound, upper_bound); + mid = (::cuda::std::min) (mid, num_items - 1); + + if (val < input[mid]) + { + upper_bound = mid; + } + else + { + lower_bound = mid + 1; + } + } + + return lower_bound; + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + InitWithRunOffsets(ItemT (&run_values)[RunsPerThread], RunOffsetT (&run_offsets)[RunsPerThread]) + { + // Keep the runs' items and the offsets of each run's beginning in the temporary storage + RunOffsetT thread_dst_offset = static_cast(linear_tid) * static_cast(RunsPerThread); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < RunsPerThread; i++) + { + temp_storage.runs.run_values[thread_dst_offset] = run_values[i]; + temp_storage.runs.run_offsets[thread_dst_offset] = run_offsets[i]; + thread_dst_offset++; + } + + // Ensure run offsets and run values have been written to shared memory + __syncthreads(); + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InitWithRunLengths( + ItemT (&run_values)[RunsPerThread], RunLengthT (&run_lengths)[RunsPerThread], TotalDecodedSizeT& total_decoded_size) + { + // Compute the offset for the beginning of each run + DecodedOffsetT run_offsets[RunsPerThread]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < RunsPerThread; i++) + { + run_offsets[i] = static_cast(run_lengths[i]); + } + DecodedOffsetT decoded_size_aggregate; + RunOffsetScanT(this->temp_storage.offset_scan).ExclusiveSum(run_offsets, run_offsets, decoded_size_aggregate); + total_decoded_size = static_cast(decoded_size_aggregate); + + // Ensure the prefix scan's temporary storage can be reused (may be superfluous, but depends on scan implementation) + __syncthreads(); + + InitWithRunOffsets(run_values, run_offsets); + } + +public: + /** + * \brief Run-length decodes the runs previously passed via a call to Init(...) and returns the run-length decoded + * items in a blocked arrangement to \p decoded_items. If the number of run-length decoded items exceeds the + * run-length decode buffer (i.e., `DecodedItemsPerThread * BLOCK_THREADS`), only the items that fit within + * the buffer are returned. Subsequent calls to `RunLengthDecode` adjusting \p from_decoded_offset can be + * used to retrieve the remaining run-length decoded items. Calling __syncthreads() between any two calls to + * `RunLengthDecode` is not required. + * \p item_offsets can be used to retrieve each run-length decoded item's relative index within its run. E.g., the + * run-length encoded array of `3, 1, 4` with the respective run lengths of `2, 1, 3` would yield the run-length + * decoded array of `3, 3, 1, 4, 4, 4` with the relative offsets of `0, 1, 0, 0, 1, 2`. + * \smemreuse + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * \param[out] decoded_items The run-length decoded items to be returned in a blocked arrangement + * \param[out] item_offsets The run-length decoded items' relative offset within the run they belong to + * \param[in] from_decoded_offset If invoked with from_decoded_offset that is larger than total_decoded_size results + * in undefined behavior. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void RunLengthDecode( + ItemT (&decoded_items)[DecodedItemsPerThread], + RelativeOffsetT (&item_offsets)[DecodedItemsPerThread], + DecodedOffsetT from_decoded_offset = 0) + { + // The (global) offset of the first item decoded by this thread + DecodedOffsetT thread_decoded_offset = from_decoded_offset + linear_tid * DecodedItemsPerThread; + + // The run that the first decoded item of this thread belongs to + // If this thread's is already beyond the total decoded size, it will be assigned to the + // last run + RunOffsetT assigned_run = + StaticUpperBound(temp_storage.runs.run_offsets, BLOCK_RUNS, thread_decoded_offset) + - static_cast(1U); + + DecodedOffsetT assigned_run_begin = temp_storage.runs.run_offsets[assigned_run]; + + // If this thread is getting assigned the last run, we make sure it will not fetch any other run after this + DecodedOffsetT assigned_run_end = + (assigned_run == BLOCK_RUNS - 1) + ? thread_decoded_offset + DecodedItemsPerThread + : temp_storage.runs.run_offsets[assigned_run + 1]; + + ItemT val = temp_storage.runs.run_values[assigned_run]; + + _CCCL_PRAGMA_UNROLL_FULL() + for (DecodedOffsetT i = 0; i < DecodedItemsPerThread; i++) + { + decoded_items[i] = val; + item_offsets[i] = thread_decoded_offset - assigned_run_begin; + + // A thread only needs to fetch the next run if this was not the last loop iteration + const bool is_final_loop_iteration = (i + 1 >= DecodedItemsPerThread); + if (!is_final_loop_iteration && (thread_decoded_offset == assigned_run_end - 1)) + { + // We make sure that a thread is not re-entering this conditional when being assigned to the last run already by + // extending the last run's length to all the thread's item + assigned_run++; + assigned_run_begin = temp_storage.runs.run_offsets[assigned_run]; + + // If this thread is getting assigned the last run, we make sure it will not fetch any other run after this + assigned_run_end = (assigned_run == BLOCK_RUNS - 1) + ? thread_decoded_offset + DecodedItemsPerThread + : temp_storage.runs.run_offsets[assigned_run + 1]; + val = temp_storage.runs.run_values[assigned_run]; + } + thread_decoded_offset++; + } + } + + /** + * \brief Run-length decodes the runs previously passed via a call to Init(...) and returns the run-length decoded + * items in a blocked arrangement to `decoded_items`. If the number of run-length decoded items exceeds the + * run-length decode buffer (i.e., `DecodedItemsPerThread * BLOCK_THREADS`), only the items that fit within + * the buffer are returned. Subsequent calls to `RunLengthDecode` adjusting `from_decoded_offset` can be + * used to retrieve the remaining run-length decoded items. Calling __syncthreads() between any two calls to + * `RunLengthDecode` is not required. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * \param[out] decoded_items The run-length decoded items to be returned in a blocked arrangement + * \param[in] from_decoded_offset If invoked with from_decoded_offset that is larger than total_decoded_size results + * in undefined behavior. + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void + RunLengthDecode(ItemT (&decoded_items)[DecodedItemsPerThread], DecodedOffsetT from_decoded_offset = 0) + { + DecodedOffsetT item_offsets[DecodedItemsPerThread]; + RunLengthDecode(decoded_items, item_offsets, from_decoded_offset); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_scan.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_scan.cuh new file mode 100644 index 00000000..43b06f82 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_scan.cuh @@ -0,0 +1,2277 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! The cub::BlockScan class provides :ref:`collective ` methods for computing a parallel prefix +//! sum/scan of items partitioned across a CUDA thread block. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Algorithmic variants + ******************************************************************************/ + +//! @brief BlockScanAlgorithm enumerates alternative algorithms for cub::BlockScan to compute a +//! parallel prefix scan across a CUDA thread block. +enum BlockScanAlgorithm +{ + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! An efficient "raking reduce-then-scan" prefix scan algorithm. Execution is comprised of five phases: + //! + //! #. Upsweep sequential reduction in registers (if threads contribute more than one input each). + //! Each thread then places the partial reduction of its item(s) into shared memory. + //! #. Upsweep sequential reduction in shared memory. + //! Threads within a single warp rake across segments of shared partial reductions. + //! #. A warp-synchronous Kogge-Stone style exclusive scan within the raking warp. + //! #. Downsweep sequential exclusive scan in shared memory. + //! Threads within a single warp rake across segments of shared partial reductions, + //! seeded with the warp-scan output. + //! #. Downsweep sequential scan in registers (if threads contribute more than one input), + //! seeded with the raking scan output. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - Although this variant may suffer longer turnaround latencies when the + //! GPU is under-occupied, it can often provide higher overall throughput + //! across the GPU when suitably occupied. + //! + //! @endrst + BLOCK_SCAN_RAKING, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! Similar to cub::BLOCK_SCAN_RAKING, but with fewer shared memory reads at the expense of higher + //! register pressure. Raking threads preserve their "upsweep" segment of values in registers while performing + //! warp-synchronous scan, allowing the "downsweep" not to re-read them from shared memory. + //! + //! @endrst + BLOCK_SCAN_RAKING_MEMOIZE, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A quick "tiled warpscans" prefix scan algorithm. Execution is comprised of four phases: + //! #. Upsweep sequential reduction in registers (if threads contribute more than one input each). + //! Each thread then places the partial reduction of its item(s) into shared memory. + //! #. Compute a shallow, but inefficient warp-synchronous Kogge-Stone style scan within each warp. + //! #. A propagation phase where the warp scan outputs in each warp are updated with the aggregate + //! from each preceding warp. + //! #. Downsweep sequential scan in registers (if threads contribute more than one input), + //! seeded with the raking scan output. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - Although this variant may suffer lower overall throughput across the + //! GPU because due to a heavy reliance on inefficient warpscans, it can + //! often provide lower turnaround latencies when the GPU is under-occupied. + //! + //! @endrst + BLOCK_SCAN_WARP_SCANS, +}; + +#if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED) +namespace detail +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(BlockScanAlgorithm algo) noexcept +{ + switch (algo) + { + case BLOCK_SCAN_RAKING: + return "BLOCK_SCAN_RAKING"; + case BLOCK_SCAN_RAKING_MEMOIZE: + return "BLOCK_SCAN_RAKING_MEMOIZE"; + case BLOCK_SCAN_WARP_SCANS: + return "BLOCK_SCAN_WARP_SCANS"; + } + return ""; +} +} // namespace detail + +inline ::std::ostream& operator<<(::std::ostream& os, BlockScanAlgorithm algo) +{ + return os << CUB_NS_QUALIFIER::detail::to_string(algo); +} +#endif // _CCCL_HOSTED() && !_CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END + +#if __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) +template <::cuda::std::same_as CharT> +struct std::formatter : formatter +{ + template + auto format(const CUB_NS_QUALIFIER::BlockScanAlgorithm& algo, FmtCtx& ctx) const + { + return formatter::format(CUB_NS_QUALIFIER::detail::to_string(algo), ctx); + } +}; +#endif // __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) + +CUB_NAMESPACE_BEGIN + +//! @rst +//! The BlockScan class provides :ref:`collective ` methods for computing a parallel prefix +//! sum/scan of items partitioned across a CUDA thread block. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - Given a list of input elements and a binary reduction operator, a +//! `prefix scan `_ produces an output list where each element is computed +//! to be the reduction of the elements occurring earlier in the input list. *Prefix sum* connotes a prefix scan with +//! the addition operator. The term *inclusive indicates* that the *i*\ :sup:`th` output reduction incorporates +//! the *i*\ :sup:`th` input. The term *exclusive* indicates the *i*\ :sup:`th` input is not incorporated into +//! the *i*\ :sup:`th` output reduction. +//! - @rowmajor +//! - BlockScan can be optionally specialized by algorithm to accommodate different workload profiles: +//! +//! #. :cpp:enumerator:`cub::BLOCK_SCAN_RAKING`: +//! An efficient (high throughput) "raking reduce-then-scan" prefix scan algorithm. +//! #. :cpp:enumerator:`cub::BLOCK_SCAN_RAKING_MEMOIZE`: +//! Similar to cub::BLOCK_SCAN_RAKING, but having higher throughput at the expense of additional +//! register pressure for intermediate storage. +//! #. :cpp:enumerator:`cub::BLOCK_SCAN_WARP_SCANS`: +//! A quick (low latency) "tiled warpscans" prefix scan algorithm. +//! +//! Performance Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - @granularity +//! - Uses special instructions when applicable (e.g., warp ``SHFL``) +//! - Uses synchronization-free communication between warp lanes when applicable +//! - Invokes a minimal number of minimal block-wide synchronization barriers (only +//! one or two depending on algorithm selection) +//! - Incurs zero bank conflicts for most types +//! - Computation is slightly more efficient (i.e., having lower instruction overhead) for: +//! +//! - Prefix sum variants (vs. generic scan) +//! - @blocksize +//! +//! - See cub::BlockScanAlgorithm for performance details regarding algorithmic alternatives +//! +//! A Simple Example +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @blockcollective{BlockScan} +//! +//! The code snippet below illustrates an exclusive prefix sum of 512 integer items that +//! are partitioned in a :ref:`blocked arrangement ` across 128 threads +//! where each thread owns 4 consecutive items. +//! +//! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin exclusive-sum-array +//! :end-before: example-end exclusive-sum-array +//! +//! Suppose the set of input ``thread_data`` across the block of threads is +//! ``{[1,1,1,1], [1,1,1,1], ..., [1,1,1,1]}``. +//! The corresponding output ``thread_data`` in those threads will be +//! ``{[0,1,2,3], [4,5,6,7], ..., [508,509,510,511]}``. +//! +//! Re-using dynamically allocating shared memory +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of dynamically shared memory with +//! BlockReduce and how to re-purpose the same memory region. +//! This example can be easily adapted to the storage required by BlockScan. +//! +//! @endrst +//! +//! @tparam T +//! Data type being scanned +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam Algorithm +//! **[optional]** cub::BlockScanAlgorithm enumerator specifying the underlying algorithm to use +//! (default: cub::BLOCK_SCAN_RAKING) +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension +//! (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockScan +{ +private: + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + /** + * Ensure the template parameterization meets the requirements of the + * specified algorithm. Currently, the BLOCK_SCAN_WARP_SCANS policy + * cannot be used with thread block sizes not a multiple of the + * architectural warp size. + */ + static constexpr BlockScanAlgorithm SAFE_ALGORITHM = + ((Algorithm == BLOCK_SCAN_WARP_SCANS) && (BLOCK_THREADS % detail::warp_threads != 0)) + ? BLOCK_SCAN_RAKING + : Algorithm; + + using WarpScans = detail::BlockScanWarpScans; + using Raking = + detail::BlockScanRaking; + + /// Define the delegate type for the desired algorithm + using InternalBlockScan = ::cuda::std::_If; + + /// Shared memory storage layout type for BlockScan + using _TempStorage = typename InternalBlockScan::TempStorage; + + /// Shared storage reference + _TempStorage& temp_storage; + + /// Linear thread-id + unsigned int linear_tid; + + /// Internal storage allocator + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + +public: + /// @smemstorage{BlockScan} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + //! @brief Collective constructor using a private static allocation of shared memory as temporary storage. + //! + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + _CCCL_DEVICE _CCCL_FORCEINLINE BlockScan() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @brief Collective constructor using the specified memory allocation as temporary storage. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @param[in] temp_storage + * Reference to memory allocation having layout type TempStorage + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockScan(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Exclusive prefix sum operations + //! @{ + + //! @rst + //! Computes an exclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes one input element. The value of 0 is applied as the initial value, and is assigned + //! to ``output`` in *thread*\ :sub:`0`. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @identityzero + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an exclusive prefix sum of 128 integer items that + //! are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-single + //! :end-before: example-end exclusive-sum-single + //! + //! Suppose the set of input ``thread_data`` across the block of threads is ``1, 1, ..., 1``. + //! The corresponding output ``thread_data`` in those threads will be ``0, 1, ..., 127``. + //! + //! @endrst + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output) + { + T initial_value{}; + + ExclusiveScan(input, output, initial_value, ::cuda::std::plus<>{}); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes one input element. + //! The value of 0 is applied as the initial value, and is assigned to ``output`` in *thread*\ :sub:`0`. + //! Also provides every thread with the block-wide ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @identityzero + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an exclusive prefix sum of 128 integer items that + //! are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-aggregate + //! :end-before: example-end exclusive-sum-aggregate + //! + //! Suppose the set of input ``thread_data`` across the block of threads is ``1, 1, ..., 1``. + //! The corresponding output ``thread_data`` in those threads will be ``0, 1, ..., 127``. + //! Furthermore the value ``128`` will be stored in ``block_aggregate`` for all threads. + //! + //! @endrst + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[out] block_aggregate + //! block-wide aggregate reduction of input items + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output, T& block_aggregate) + { + T initial_value{}; + + ExclusiveScan(input, output, initial_value, ::cuda::std::plus<>{}, block_aggregate); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes one input element. Instead of using 0 as the block-wide prefix, the call-back functor + //! ``block_prefix_callback_op`` is invoked by the first warp in the block, and the value returned by + //! *lane*\ :sub:`0` in that warp is used as the "seed" value that logically prefixes the thread block's + //! scan inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @identityzero + //! - The ``block_prefix_callback_op`` functor must implement a member function + //! ``T operator()(T block_aggregate)``. The functor will be invoked by the first warp of threads in the block, + //! however only the return value from *lane*\ :sub:`0` is applied as the block-wide prefix. Can be stateful. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a single thread block that progressively + //! computes an exclusive prefix sum over multiple "tiles" of input using a + //! prefix functor to maintain a running total between block-wide scans. Each tile consists + //! of 128 integer items that are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin block-prefix-callback-op + //! :end-before: example-end block-prefix-callback-op + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-single-prefix-callback + //! :end-before: example-end exclusive-sum-single-prefix-callback + //! + //! Suppose the input ``d_data`` is ``1, 1, 1, 1, 1, 1, 1, 1, ...``. + //! The corresponding output for the first segment will be ``0, 1, ..., 127``. + //! The output for the second segment will be ``128, 129, ..., 255``. + //! + //! @endrst + //! + //! @tparam BlockPrefixCallbackOp + //! **[inferred]** Call-back functor type having member `T operator()(T block_aggregate)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in,out] block_prefix_callback_op + //! @rst + //! *warp*\ :sub:`0` only call-back functor for specifying a block-wide prefix to be applied to + //! the logical input sequence. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output, BlockPrefixCallbackOp& block_prefix_callback_op) + { + ExclusiveScan(input, output, ::cuda::std::plus<>{}, block_prefix_callback_op); + } + + //! @} + //! @name Exclusive prefix sum operations (multiple data per thread) + //! @{ + + //! @rst + //! Computes an exclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes an array of consecutive input elements. + //! The value of 0 is applied as the initial value, and is assigned to ``output[0]`` in *thread*\ :sub:`0`. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @identityzero + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an exclusive prefix sum of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-array + //! :end-before: example-end exclusive-sum-array + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [1,1,1,1], [1,1,1,1], ..., [1,1,1,1] }``. + //! The corresponding output ``thread_data`` in those threads will be + //! ``{ [0,1,2,3], [4,5,6,7], ..., [508,509,510,511] }``. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD]) + { + T initial_value{}; + + ExclusiveScan(input, output, initial_value, ::cuda::std::plus<>{}); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes an array of consecutive input elements. + //! The value of 0 is applied as the initial value, and is assigned to ``output[0]`` in *thread*\ :sub:`0`. + //! Also provides every thread with the block-wide ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @identityzero + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an exclusive prefix sum of 512 integer items that are partitioned in + //! a :ref:`blocked arrangement ` across 128 threads where each thread owns + //! 4 consecutive items. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-array-aggregate + //! :end-before: example-end exclusive-sum-array-aggregate + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [1,1,1,1], [1,1,1,1], ..., [1,1,1,1] }``. + //! The corresponding output ``thread_data`` in those threads will be + //! ``{ [0,1,2,3], [4,5,6,7], ..., [508,509,510,511] }``. + //! Furthermore the value ``512`` will be stored in ``block_aggregate`` for all threads. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[out] block_aggregate + //! block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T& block_aggregate) + { + // Reduce consecutive thread items in registers + T initial_value{}; + + ExclusiveScan(input, output, initial_value, ::cuda::std::plus<>{}, block_aggregate); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes an array of consecutive input elements. + //! Instead of using 0 as the block-wide prefix, the call-back functor ``block_prefix_callback_op`` is invoked by + //! the first warp in the block, and the value returned by *lane*\ :sub:`0` in that warp is used as the "seed" + //! value that logically prefixes the thread block's scan inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @identityzero + //! - The ``block_prefix_callback_op`` functor must implement a member function ``T operator()(T block_aggregate)``. + //! The functor will be invoked by the first warp of threads in the block, however only the return value from + //! *lane*\ :sub:`0` is applied as the block-wide prefix. Can be stateful. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a single thread block that progressively + //! computes an exclusive prefix sum over multiple "tiles" of input using a + //! prefix functor to maintain a running total between block-wide scans. Each tile consists + //! of 512 integer items that are partitioned in a :ref:`blocked arrangement ` + //! across 128 threads where each thread owns 4 consecutive items. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin block-prefix-callback-op + //! :end-before: example-end block-prefix-callback-op + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-prefix-callback + //! :end-before: example-end exclusive-sum-prefix-callback + //! + //! Suppose the input ``d_data`` is ``1, 1, 1, 1, 1, 1, 1, 1, ...``. + //! The corresponding output for the first segment will be ``0, 1, 2, 3, ..., 510, 511``. + //! The output for the second segment will be ``512, 513, 514, 515, ..., 1022, 1023``. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam BlockPrefixCallbackOp + //! **[inferred]** Call-back functor type having member + //! `T operator()(T block_aggregate)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in,out] block_prefix_callback_op + //! @rst + //! *warp*\ :sub:`0` only call-back functor for specifying a block-wide prefix to be applied to + //! the logical input sequence. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum( + T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], BlockPrefixCallbackOp& block_prefix_callback_op) + { + ExclusiveScan(input, output, ::cuda::std::plus<>{}, block_prefix_callback_op); + } + + //! @} + //! @name Exclusive prefix scan operations + //! @{ + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes one input element. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an exclusive prefix max scan of 128 integer items that + //! are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-single + //! :end-before: example-end exclusive-scan-single + //! + //! Suppose the set of input ``thread_data`` across the block of threads is ``0, -1, 2, -3, ..., 126, -127``. + //! The corresponding output ``thread_data`` in those threads will be ``INT_MIN, 0, 0, 2, ..., 124, 126``. + //! + //! @endrst + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in] initial_value + //! @rst + //! Initial value to seed the exclusive scan (and is assigned to `output[0]` in *thread*\ :sub:`0`) + //! @endrst + //! + //! @param[in] scan_op + //! Binary scan functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, T initial_value, ScanOp scan_op) + { + InternalBlockScan(temp_storage).ExclusiveScan(input, output, initial_value, scan_op); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes one input element. + //! Also provides every thread with the block-wide ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an exclusive prefix max scan of 128 integer items that + //! are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-aggregate + //! :end-before: example-end exclusive-scan-aggregate + //! + //! Suppose the set of input ``thread_data`` across the block of threads is ``0, -1, 2, -3, ..., 126, -127``. + //! The corresponding output ``thread_data`` in those threads will be ``INT_MIN, 0, 0, 2, ..., 124, 126``. + //! Furthermore the value ``126`` will be stored in ``block_aggregate`` for all threads. + //! + //! .. note:: + //! + //! ``initial_value`` is not applied to the block-wide aggregate. + //! + //! @endrst + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member ``T operator()(const T &a, const T &b)`` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to ``input``) + //! + //! @param[in] initial_value + //! @rst + //! Initial value to seed the exclusive scan (and is assigned to ``output[0]`` in *thread*\ :sub:`0`). It is not + //! taken into account for ``block_aggregate``. + //! + //! @endrst + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[out] block_aggregate + //! block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T input, T& output, T initial_value, ScanOp scan_op, T& block_aggregate) + { + InternalBlockScan(temp_storage).ExclusiveScan(input, output, initial_value, scan_op, block_aggregate); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes one input element. The call-back functor ``block_prefix_callback_op`` is invoked by + //! the first warp in the block, and the value returned by *lane*\ :sub:`0` in that warp is used as + //! the "seed" value that logically prefixes the thread block's scan inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The ``block_prefix_callback_op`` functor must implement a member function ``T operator()(T block_aggregate)``. + //! The functor will be invoked by the first warp of threads in the block, however only the return value from + //! *lane*\ :sub:`0` is applied as the block-wide prefix. Can be stateful. + //! - Supports non-commutative scan operators. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a single thread block that progressively + //! computes an exclusive prefix max scan over multiple "tiles" of input using a + //! prefix functor to maintain a running total between block-wide scans. + //! Each tile consists of 128 integer items that are partitioned across 128 threads. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // A stateful callback functor that maintains a running prefix to be applied + //! // during consecutive scan operations. + //! struct BlockPrefixCallbackOp + //! { + //! // Running prefix + //! int running_total; + //! + //! // Constructor + //! __device__ BlockPrefixCallbackOp(int running_total) : running_total(running_total) {} + //! + //! // Callback operator to be entered by the first warp of threads in the block. + //! // Thread-0 is responsible for returning a value for seeding the block-wide scan. + //! __device__ int operator()(int block_aggregate) + //! { + //! int old_prefix = running_total; + //! running_total = (block_aggregate > old_prefix) ? block_aggregate : old_prefix; + //! return old_prefix; + //! } + //! }; + //! + //! __global__ void ExampleKernel(int *d_data, int num_items, ...) + //! { + //! // Specialize BlockScan for a 1D block of 128 threads + //! using BlockScan = cub::BlockScan; + //! + //! // Allocate shared memory for BlockScan + //! __shared__ typename BlockScan::TempStorage temp_storage; + //! + //! // Initialize running total + //! BlockPrefixCallbackOp prefix_op(INT_MIN); + //! + //! // Have the block iterate over segments of items + //! for (int block_offset = 0; block_offset < num_items; block_offset += 128) + //! { + //! // Load a segment of consecutive items that are blocked across threads + //! int thread_data = d_data[block_offset + threadIdx.x]; + //! + //! // Collectively compute the block-wide exclusive prefix max scan + //! BlockScan(temp_storage).ExclusiveScan( + //! thread_data, thread_data, INT_MIN, cuda::maximum<>{}, prefix_op); + //! __syncthreads(); + //! + //! // Store scanned items to output segment + //! d_data[block_offset + threadIdx.x] = thread_data; + //! } + //! } + //! + //! Suppose the input ``d_data`` is ``0, -1, 2, -3, 4, -5, ...``. + //! The corresponding output for the first segment will be ``INT_MIN, 0, 0, 2, ..., 124, 126``. + //! The output for the second segment will be ``126, 128, 128, 130, ..., 252, 254``. + //! + //! @endrst + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam BlockPrefixCallbackOp + //! **[inferred]** Call-back functor type having member `T operator()(T block_aggregate)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[in,out] block_prefix_callback_op + //! @rst + //! *warp*\ :sub:`0` only call-back functor for specifying a block-wide prefix to be applied to + //! the logical input sequence. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T input, T& output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op) + { + InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op, block_prefix_callback_op); + } + + //! @} + //! @name Exclusive prefix scan operations (multiple data per thread) + //! @{ + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an exclusive prefix max scan of 512 integer + //! items that are partitioned in a [blocked arrangement](../index.html#sec5sec3) + //! across 128 threads where each thread owns 4 consecutive items. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-array + //! :end-before: example-end exclusive-scan-array + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,-1,2,-3], [4,-5,6,-7], ..., [508,-509,510,-511] }``. + //! The corresponding output ``thread_data`` in those threads will be + //! ``{ [INT_MIN,0,0,2], [2,4,4,6], ..., [506,508,508,510] }``. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member + //! `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] initial_value + //! @rst + //! Initial value to seed the exclusive scan (and is assigned to `output[0]` in *thread*\ :sub:`0`) + //! @endrst + //! + //! @param[in] scan_op + //! Binary scan functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op) + { + // Reduce consecutive thread items in registers + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op); + + // Exclusive scan in registers with prefix as seed + detail::ThreadScanExclusive(input, output, scan_op, thread_prefix); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. + //! Also provides every thread with the block-wide ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an exclusive prefix max scan of 512 integer items that are partitioned in + //! a :ref:`blocked arrangement ` across 128 threads where each thread owns + //! 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockScan for a 1D block of 128 threads of type int + //! using BlockScan = cub::BlockScan; + //! + //! // Allocate shared memory for BlockScan + //! __shared__ typename BlockScan::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively compute the block-wide exclusive prefix max scan + //! int block_aggregate; + //! BlockScan(temp_storage).ExclusiveScan( + //! thread_data, thread_data, INT_MIN, cuda::maximum<>{}, block_aggregate); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,-1,2,-3], [4,-5,6,-7], ..., [508,-509,510,-511] }``. + //! The corresponding output ``thread_data`` in those threads will be + //! ``{ [INT_MIN,0,0,2], [2,4,4,6], ..., [506,508,508,510] }``. + //! Furthermore the value ``510`` will be stored in ``block_aggregate`` for all threads. + //! + //! .. note:: + //! + //! ``initial_value`` is not applied to the block-wide aggregate. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] initial_value + //! @rst + //! Initial value to seed the exclusive scan (and is assigned to `output[0]` in *thread*\ :sub:`0`). It is not taken + //! into account for ``block_aggregate``. + //! @endrst + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[out] block_aggregate + //! block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan( + T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op, T& block_aggregate) + { + // Reduce consecutive thread items in registers + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op, block_aggregate); + + // Exclusive scan in registers with prefix as seed + detail::ThreadScanExclusive(input, output, scan_op, thread_prefix); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. + //! The call-back functor ``block_prefix_callback_op`` is invoked by the first warp in the block, and the value + //! returned by *lane*\ :sub:`0` in that warp is used as the "seed" value that logically prefixes the thread + //! block's scan inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The ``block_prefix_callback_op`` functor must implement a member function + //! ``T operator()(T block_aggregate)``. The functor will be invoked by the + //! first warp of threads in the block, however only the return value from + //! *lane*\ :sub:`0` is applied as the block-wide prefix. Can be stateful. + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a single thread block that progressively + //! computes an exclusive prefix max scan over multiple "tiles" of input using a + //! prefix functor to maintain a running total between block-wide scans. Each tile consists + //! of 128 integer items that are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin block-prefix-callback-max-op + //! :end-before: example-end block-prefix-callback-max-op + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-prefix-callback + //! :end-before: example-end exclusive-scan-prefix-callback + //! + //! Suppose the input ``d_data`` is ``0, -1, 2, -3, 4, -5, ...``. + //! The corresponding output for the first segment will be + //! ``INT_MIN, 0, 0, 2, 2, 4, ..., 508, 510``. + //! The output for the second segment will be + //! ``510, 512, 512, 514, 514, 516, ..., 1020, 1022``. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam BlockPrefixCallbackOp + //! **[inferred]** Call-back functor type having member `T operator()(T block_aggregate)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[in,out] block_prefix_callback_op + //! @rst + //! *warp*\ :sub:`0` only call-back functor for specifying a block-wide prefix to be applied to + //! the logical input sequence. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan( + T (&input)[ITEMS_PER_THREAD], + T (&output)[ITEMS_PER_THREAD], + ScanOp scan_op, + BlockPrefixCallbackOp& block_prefix_callback_op) + { + // Reduce consecutive thread items in registers + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_prefix_callback_op); + + // Exclusive scan in registers with prefix as seed + detail::ThreadScanExclusive(input, output, scan_op, thread_prefix); + } + + //! @} +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document no-initial-value scans + + //! @name Exclusive prefix scan operations (no initial value, single datum per thread) + //! @{ + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes one input element. + //! With no initial value, the output computed for *thread*\ :sub:`0` is undefined. + //! + //! - Supports non-commutative scan operators. + //! - @rowmajor + //! - @smemreuse + //! + //! @endrst + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, ScanOp scan_op) + { + InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes one input element. Also provides every thread with the block-wide + //! ``block_aggregate`` of all inputs. With no initial value, the output computed for + //! *thread*\ :sub:`0` is undefined. + //! + //! - Supports non-commutative scan operators. + //! - @rowmajor + //! - @smemreuse + //! + //! @endrst + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[out] block_aggregate + //! block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, ScanOp scan_op, T& block_aggregate) + { + InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op, block_aggregate); + } + + //! @} + //! @name Exclusive prefix scan operations (no initial value, multiple data per thread) + //! @{ + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. With no initial value, the + //! output computed for *thread*\ :sub:`0` is undefined. + //! + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op) + { + // Reduce consecutive thread items in registers + T thread_partial = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_partial, thread_partial, scan_op); + + // Exclusive scan in registers with prefix + detail::ThreadScanExclusive(input, output, scan_op, thread_partial, (linear_tid != 0)); + } + + //! @rst + //! Computes an exclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. Also provides every thread + //! with the block-wide ``block_aggregate`` of all inputs. + //! With no initial value, the output computed for *thread*\ :sub:`0` is undefined. + //! + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[out] block_aggregate + //! block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op, T& block_aggregate) + { + // Reduce consecutive thread items in registers + T thread_partial = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_partial, thread_partial, scan_op, block_aggregate); + + // Exclusive scan in registers with prefix + detail::ThreadScanExclusive(input, output, scan_op, thread_partial, (linear_tid != 0)); + } + + //! @} +#endif // _CCCL_DOXYGEN_INVOKED // Do not document no-initial-value scans + + //! @name Inclusive prefix sum operations + //! @{ + + //! @rst + //! Computes an inclusive block-wide prefix scan using addition (+) + //! as the scan operator. Each thread contributes one input element. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix sum of 128 integer items that + //! are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-sum-single + //! :end-before: example-end inclusive-sum-single + //! + //! Suppose the set of input ``thread_data`` across the block of threads is ``1, 1, ..., 1``. + //! The corresponding output ``thread_data`` in those threads will be ``1, 2, ..., 128``. + //! + //! @endrst + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output) + { + InclusiveScan(input, output, ::cuda::std::plus<>{}); + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes one input element. + //! Also provides every thread with the block-wide ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix sum of 128 integer items that + //! are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-sum-single-aggregate + //! :end-before: example-end inclusive-sum-single-aggregate + //! + //! Suppose the set of input ``thread_data`` across the block of threads is ``1, 1, ..., 1``. + //! The corresponding output ``thread_data`` in those threads will be ``1, 2, ..., 128``. + //! Furthermore the value ``128`` will be stored in ``block_aggregate`` for all threads. + //! + //! @endrst + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[out] block_aggregate + //! block-wide aggregate reduction of input items + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output, T& block_aggregate) + { + InclusiveScan(input, output, ::cuda::std::plus<>{}, block_aggregate); + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes one input element. Instead of using 0 as the block-wide prefix, the call-back functor + //! ``block_prefix_callback_op`` is invoked by the first warp in the block, and the value returned by + //! *lane*\ :sub:`0` in that warp is used as the "seed" value that logically prefixes the thread block's + //! scan inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The ``block_prefix_callback_op`` functor must implement a member function + //! ``T operator()(T block_aggregate)``. The functor will be invoked by the first warp of threads in the block, + //! however only the return value from *lane*\ :sub:`0` is applied as the block-wide prefix. Can be stateful. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a single thread block that progressively + //! computes an inclusive prefix sum over multiple "tiles" of input using a + //! prefix functor to maintain a running total between block-wide scans. + //! Each tile consists of 128 integer items that are partitioned across 128 threads. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // A stateful callback functor that maintains a running prefix to be applied + //! // during consecutive scan operations. + //! struct BlockPrefixCallbackOp + //! { + //! // Running prefix + //! int running_total; + //! + //! // Constructor + //! __device__ BlockPrefixCallbackOp(int running_total) : running_total(running_total) {} + //! + //! // Callback operator to be entered by the first warp of threads in the block. + //! // Thread-0 is responsible for returning a value for seeding the block-wide scan. + //! __device__ int operator()(int block_aggregate) + //! { + //! int old_prefix = running_total; + //! running_total += block_aggregate; + //! return old_prefix; + //! } + //! }; + //! + //! __global__ void ExampleKernel(int *d_data, int num_items, ...) + //! { + //! // Specialize BlockScan for a 1D block of 128 threads + //! using BlockScan = cub::BlockScan; + //! + //! // Allocate shared memory for BlockScan + //! __shared__ typename BlockScan::TempStorage temp_storage; + //! + //! // Initialize running total + //! BlockPrefixCallbackOp prefix_op(0); + //! + //! // Have the block iterate over segments of items + //! for (int block_offset = 0; block_offset < num_items; block_offset += 128) + //! { + //! // Load a segment of consecutive items that are blocked across threads + //! int thread_data = d_data[block_offset + threadIdx.x]; + //! + //! // Collectively compute the block-wide inclusive prefix sum + //! BlockScan(temp_storage).InclusiveSum( + //! thread_data, thread_data, prefix_op); + //! __syncthreads(); + //! + //! // Store scanned items to output segment + //! d_data[block_offset + threadIdx.x] = thread_data; + //! } + //! + //! Suppose the input ``d_data`` is ``1, 1, 1, 1, 1, 1, 1, 1, ...``. + //! The corresponding output for the first segment will be ``1, 2, ..., 128``. + //! The output for the second segment will be ``129, 130, ..., 256``. + //! + //! @endrst + //! + //! @tparam BlockPrefixCallbackOp + //! **[inferred]** Call-back functor type having member `T operator()(T block_aggregate)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in,out] block_prefix_callback_op + //! @rst + //! *warp*\ :sub:`0` only call-back functor for specifying a block-wide prefix to be applied + //! to the logical input sequence. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output, BlockPrefixCallbackOp& block_prefix_callback_op) + { + InclusiveScan(input, output, ::cuda::std::plus<>{}, block_prefix_callback_op); + } + + //! @} + //! @name Inclusive prefix sum operations (multiple data per thread) + //! @{ + + //! @rst + //! Computes an inclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes an array of consecutive input elements. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix sum of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-sum-array + //! :end-before: example-end inclusive-sum-array + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [1,1,1,1], [1,1,1,1], ..., [1,1,1,1] }``. The corresponding output + //! ``thread_data`` in those threads will be ``{ [1,2,3,4], [5,6,7,8], ..., [509,510,511,512] }``. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD]) + { + if constexpr (ITEMS_PER_THREAD == 1) + { + InclusiveSum(input[0], output[0]); + } + else + { + // Reduce consecutive thread items in registers + ::cuda::std::plus<> scan_op; + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveSum(thread_prefix, thread_prefix); + + // Inclusive scan in registers with prefix as seed + detail::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0)); + } + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes an array of consecutive input elements. + //! Also provides every thread with the block-wide ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix sum of 512 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-sum-array-aggregate + //! :end-before: example-end inclusive-sum-array-aggregate + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [1,1,1,1], [1,1,1,1], ..., [1,1,1,1] }``. The + //! corresponding output ``thread_data`` in those threads will be + //! ``{ [1,2,3,4], [5,6,7,8], ..., [509,510,511,512] }``. + //! Furthermore the value ``512`` will be stored in ``block_aggregate`` for all threads. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[out] block_aggregate + //! block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + InclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T& block_aggregate) + { + if constexpr (ITEMS_PER_THREAD == 1) + { + InclusiveSum(input[0], output[0], block_aggregate); + } + else + { + // Reduce consecutive thread items in registers + ::cuda::std::plus<> scan_op; + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveSum(thread_prefix, thread_prefix, block_aggregate); + + // Inclusive scan in registers with prefix as seed + detail::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0)); + } + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using addition (+) as the scan operator. + //! Each thread contributes an array of consecutive input elements. + //! Instead of using 0 as the block-wide prefix, the call-back functor ``block_prefix_callback_op`` is invoked by + //! the first warp in the block, and the value returned by *lane*\ :sub:`0` in that warp is used as the "seed" + //! value that logically prefixes the thread block's scan inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The ``block_prefix_callback_op`` functor must implement a member function + //! ``T operator()(T block_aggregate)``. The functor will be invoked by the first warp of threads in the block, + //! however only the return value from *lane*\ :sub:`0` is applied as the block-wide prefix. Can be stateful. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a single thread block that progressively + //! computes an inclusive prefix sum over multiple "tiles" of input using a + //! prefix functor to maintain a running total between block-wide scans. Each tile consists + //! of 512 integer items that are partitioned in a :ref:`blocked arrangement ` + //! across 128 threads where each thread owns 4 consecutive items. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin block-prefix-callback-op + //! :end-before: example-end block-prefix-callback-op + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-prefix-callback + //! :end-before: example-end inclusive-scan-prefix-callback + //! + //! Suppose the input ``d_data`` is ``1, 1, 1, 1, 1, 1, 1, 1, ...``. + //! The corresponding output for the first segment will be + //! ``1, 2, 3, 4, ..., 511, 512``. The output for the second segment will be + //! ``513, 514, 515, 516, ..., 1023, 1024``. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam BlockPrefixCallbackOp + //! **[inferred]** Call-back functor type having member `T operator()(T block_aggregate)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in,out] block_prefix_callback_op + //! @rst + //! *warp*\ :sub:`0` only call-back functor for specifying a block-wide prefix to be applied to the + //! logical input sequence. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum( + T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], BlockPrefixCallbackOp& block_prefix_callback_op) + { + if constexpr (ITEMS_PER_THREAD == 1) + { + InclusiveSum(input[0], output[0], block_prefix_callback_op); + } + else + { + // Reduce consecutive thread items in registers + ::cuda::std::plus<> scan_op; + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveSum(thread_prefix, thread_prefix, block_prefix_callback_op); + + // Inclusive scan in registers with prefix as seed + detail::ThreadScanInclusive(input, output, scan_op, thread_prefix); + } + } + + //! @} + //! @name Inclusive prefix scan operations + //! @{ + + //! @rst + //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes one input element. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix max scan of 128 integer items that + //! are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-single + //! :end-before: example-end inclusive-scan-single + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``0, -1, 2, -3, ..., 126, -127``. The corresponding output ``thread_data`` + //! in those threads will be ``0, 0, 2, 2, ..., 126, 126``. + //! + //! @endrst + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& output, ScanOp scan_op) + { + InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op); + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes one input element. Also provides every thread with the block-wide + //! ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix max scan of 128 + //! integer items that are partitioned across 128 threads. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockScan for a 1D block of 128 threads of type int + //! using BlockScan = cub::BlockScan; + //! + //! // Allocate shared memory for BlockScan + //! __shared__ typename BlockScan::TempStorage temp_storage; + //! + //! // Obtain input item for each thread + //! int thread_data; + //! ... + //! + //! // Collectively compute the block-wide inclusive prefix max scan + //! int block_aggregate; + //! BlockScan(temp_storage).InclusiveScan(thread_data, thread_data, cuda::maximum<>{}, block_aggregate); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``0, -1, 2, -3, ..., 126, -127``. The corresponding output ``thread_data`` + //! in those threads will be ``0, 0, 2, 2, ..., 126, 126``. Furthermore the value + //! ``126`` will be stored in ``block_aggregate`` for all threads. + //! + //! @endrst + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[out] block_aggregate + //! Block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& output, ScanOp scan_op, T& block_aggregate) + { + InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op, block_aggregate); + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes one input element. The call-back functor ``block_prefix_callback_op`` + //! is invoked by the first warp in the block, and the value returned by *lane*\ :sub:`0` in that warp is used as + //! the "seed" value that logically prefixes the thread block's scan inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The ``block_prefix_callback_op`` functor must implement a member function + //! ``T operator()(T block_aggregate)``. The functor's input parameter + //! The functor will be invoked by the first warp of threads in the block, + //! however only the return value from *lane*\ :sub:`0` is applied + //! as the block-wide prefix. Can be stateful. + //! - Supports non-commutative scan operators. + //! - @rowmajor + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a single thread block that progressively + //! computes an inclusive prefix max scan over multiple "tiles" of input using a + //! prefix functor to maintain a running total between block-wide scans. Each tile consists + //! of 128 integer items that are partitioned across 128 threads. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin block-prefix-callback-max-op + //! :end-before: example-end block-prefix-callback-max-op + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-prefix-callback-max + //! :end-before: example-end inclusive-scan-prefix-callback-max + //! + //! Suppose the input ``d_data`` is ``0, -1, 2, -3, 4, -5, ...``. + //! The corresponding output for the first segment will be + //! ``0, 0, 2, 2, ..., 126, 126``. The output for the second segment + //! will be ``128, 128, 130, 130, ..., 254, 254``. + //! + //! @endrst + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam BlockPrefixCallbackOp + //! **[inferred]** Call-back functor type having member `T operator()(T block_aggregate)` + //! + //! @param[in] input + //! Calling thread's input item + //! + //! @param[out] output + //! Calling thread's output item (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[in,out] block_prefix_callback_op + //! @rst + //! *warp*\ :sub:`0` only call-back functor for specifying a block-wide prefix to be applied to + //! the logical input sequence. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + InclusiveScan(T input, T& output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op) + { + InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op, block_prefix_callback_op); + } + + //! @} + //! @name Inclusive prefix scan operations (multiple data per thread) + //! @{ + + //! @rst + //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix max scan of 512 integer items that + //! are partitioned in a [blocked arrangement](../index.html#sec5sec3) across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. literalinclude:: ../../../cub/examples/block/example_block_scan.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-array + //! :end-before: example-end inclusive-scan-array + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,-1,2,-3], [4,-5,6,-7], ..., [508,-509,510,-511] }``. + //! The corresponding output ``thread_data`` in those threads will be + //! ``{ [0,0,2,2], [4,4,6,6], ..., [508,508,510,510] }``. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op) + { + if constexpr (ITEMS_PER_THREAD == 1) + { + InclusiveScan(input[0], output[0], scan_op); + } + else + { + // Reduce consecutive thread items in registers + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_prefix, thread_prefix, scan_op); + + // Inclusive scan in registers with prefix as seed (first thread does not seed) + detail::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0)); + } + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix max scan of 128 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 64 threads + //! where each thread owns 2 consecutive items. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_scan_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-array-init-value + //! :end-before: example-end inclusive-scan-array-init-value + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] initial_value + //! Initial value to seed the inclusive scan (uniform across block) + //! + //! @param[in] scan_op + //! Binary scan functor + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op) + { + // Reduce consecutive thread items in registers + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op); + + // Exclusive scan in registers with prefix as seed + detail::ThreadScanInclusive(input, output, scan_op, thread_prefix); + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. Also provides every thread + //! with the block-wide ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix max scan of 512 integer items that + //! are partitioned in a [blocked arrangement](../index.html#sec5sec3) across 128 threads + //! where each thread owns 4 consecutive items. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(...) + //! { + //! // Specialize BlockScan for a 1D block of 128 threads of type int + //! using BlockScan = cub::BlockScan; + //! + //! // Allocate shared memory for BlockScan + //! __shared__ typename BlockScan::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Collectively compute the block-wide inclusive prefix max scan + //! int block_aggregate; + //! BlockScan(temp_storage).InclusiveScan(thread_data, thread_data, cuda::maximum<>{}, block_aggregate); + //! + //! Suppose the set of input ``thread_data`` across the block of threads is + //! ``{ [0,-1,2,-3], [4,-5,6,-7], ..., [508,-509,510,-511] }``. + //! The corresponding output ``thread_data`` in those threads will be + //! ``{ [0,0,2,2], [4,4,6,6], ..., [508,508,510,510] }``. + //! Furthermore the value ``510`` will be stored in ``block_aggregate`` for all threads. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[out] block_aggregate + //! Block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op, T& block_aggregate) + { + if (ITEMS_PER_THREAD == 1) + { + InclusiveScan(input[0], output[0], scan_op, block_aggregate); + } + else + { + // Reduce consecutive thread items in registers + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan (with no initial value) + ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_aggregate); + + // Inclusive scan in registers with prefix as seed (first thread does not seed) + detail::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0)); + } + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. Also provides every thread + //! with the block-wide ``block_aggregate`` of all inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates an inclusive prefix max scan of 128 integer items that + //! are partitioned in a :ref:`blocked arrangement ` across 64 threads + //! where each thread owns 2 consecutive items. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_block_scan_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-array-aggregate-init-value + //! :end-before: example-end inclusive-scan-array-aggregate-init-value + //! + //! The value ``126`` will be stored in ``block_aggregate`` for all threads. + //! + //! .. note:: + //! + //! ``initial_value`` is not applied to the block-wide aggregate. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] initial_value + //! Initial value to seed the inclusive scan (uniform across block). It is not taken + //! into account for ``block_aggregate``. + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[out] block_aggregate + //! Block-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan( + T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op, T& block_aggregate) + { + // Reduce consecutive thread items in registers + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op, block_aggregate); + + // Exclusive scan in registers with prefix as seed + detail::ThreadScanInclusive(input, output, scan_op, thread_prefix); + } + + //! @rst + //! Computes an inclusive block-wide prefix scan using the specified binary ``scan_op`` functor. + //! Each thread contributes an array of consecutive input elements. + //! The call-back functor ``block_prefix_callback_op`` is invoked by the first warp in the block, + //! and the value returned by *lane*\ :sub:`0` in that warp is used as the "seed" value that logically prefixes the + //! thread block's scan inputs. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The ``block_prefix_callback_op`` functor must implement a member function ``T operator()(T block_aggregate)``. + //! The functor will be invoked by the first warp of threads in the block, however only the return value + //! from *lane*\ :sub:`0` is applied as the block-wide prefix. Can be stateful. + //! - Supports non-commutative scan operators. + //! - @blocked + //! - @granularity + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates a single thread block that progressively + //! computes an inclusive prefix max scan over multiple "tiles" of input using a + //! prefix functor to maintain a running total between block-wide scans. Each tile consists + //! of 128 integer items that are partitioned across 128 threads. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // A stateful callback functor that maintains a running prefix to be applied + //! // during consecutive scan operations. + //! struct BlockPrefixCallbackOp + //! { + //! // Running prefix + //! int running_total; + //! + //! // Constructor + //! __device__ BlockPrefixCallbackOp(int running_total) : running_total(running_total) {} + //! + //! // Callback operator to be entered by the first warp of threads in the block. + //! // Thread-0 is responsible for returning a value for seeding the block-wide scan. + //! __device__ int operator()(int block_aggregate) + //! { + //! int old_prefix = running_total; + //! running_total = (block_aggregate > old_prefix) ? block_aggregate : old_prefix; + //! return old_prefix; + //! } + //! }; + //! + //! __global__ void ExampleKernel(int *d_data, int num_items, ...) + //! { + //! // Specialize BlockLoad, BlockStore, and BlockScan for a 1D block of 128 threads, 4 ints per thread + //! using BlockLoad = cub::BlockLoad ; + //! using BlockStore = cub::BlockStore ; + //! using BlockScan = cub::BlockScan ; + //! + //! // Allocate aliased shared memory for BlockLoad, BlockStore, and BlockScan + //! __shared__ union { + //! typename BlockLoad::TempStorage load; + //! typename BlockScan::TempStorage scan; + //! typename BlockStore::TempStorage store; + //! } temp_storage; + //! + //! // Initialize running total + //! BlockPrefixCallbackOp prefix_op(0); + //! + //! // Have the block iterate over segments of items + //! for (int block_offset = 0; block_offset < num_items; block_offset += 128 * 4) + //! { + //! // Load a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! BlockLoad(temp_storage.load).Load(d_data + block_offset, thread_data); + //! __syncthreads(); + //! + //! // Collectively compute the block-wide inclusive prefix max scan + //! BlockScan(temp_storage.scan).InclusiveScan( + //! thread_data, thread_data, cuda::maximum<>{}, prefix_op); + //! __syncthreads(); + //! + //! // Store scanned items to output segment + //! BlockStore(temp_storage.store).Store(d_data + block_offset, thread_data); + //! __syncthreads(); + //! } + //! + //! Suppose the input ``d_data`` is ``0, -1, 2, -3, 4, -5, ...``. + //! The corresponding output for the first segment will be + //! ``0, 0, 2, 2, 4, 4, ..., 510, 510``. The output for the second + //! segment will be ``512, 512, 514, 514, 516, 516, ..., 1022, 1022``. + //! + //! @endrst + //! + //! @tparam ITEMS_PER_THREAD + //! **[inferred]** The number of consecutive items partitioned onto each thread. + //! + //! @tparam ScanOp + //! **[inferred]** Binary scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam BlockPrefixCallbackOp + //! **[inferred]** Call-back functor type having member `T operator()(T block_aggregate)` + //! + //! @param[in] input + //! Calling thread's input items + //! + //! @param[out] output + //! Calling thread's output items (may be aliased to `input`) + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[in,out] block_prefix_callback_op + //! @rst + //! *warp*\ :sub:`0` only call-back functor for specifying a block-wide prefix to be applied to + //! the logical input sequence. + //! @endrst + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan( + T (&input)[ITEMS_PER_THREAD], + T (&output)[ITEMS_PER_THREAD], + ScanOp scan_op, + BlockPrefixCallbackOp& block_prefix_callback_op) + { + if (ITEMS_PER_THREAD == 1) + { + InclusiveScan(input[0], output[0], scan_op, block_prefix_callback_op); + } + else + { + // Reduce consecutive thread items in registers + T thread_prefix = cub::ThreadReduce(input, scan_op); + + // Exclusive thread block-scan + ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_prefix_callback_op); + + // Inclusive scan in registers with prefix as seed + detail::ThreadScanInclusive(input, output, scan_op, thread_prefix); + } + } + + //! @} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_store.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_store.cuh new file mode 100644 index 00000000..59113bfc --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_store.cuh @@ -0,0 +1,960 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! Operations for writing linear segments of data from the CUDA thread block + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @name Blocked arrangement I/O (direct) +//! @{ + +//! @rst +//! Store a blocked arrangement of items across a thread block into a linear segment of items +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @blocked +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to store. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam OutputIteratorT +//! **[inferred]** The random-access iterator type for output @iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread +//! (e.g., ``(threadIdx.y * blockDim.x) + linear_tid`` for 2D thread blocks) +//! +//! @param[in] block_itr +//! The thread block's base output iterator for storing to +//! +//! @param[in] items +//! Data to store +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +StoreDirectBlocked(int linear_tid, OutputIteratorT block_itr, T (&items)[ItemsPerThread]) +{ + OutputIteratorT thread_itr = block_itr + (linear_tid * ItemsPerThread); // NOLINT(bugprone-misplaced-widening-cast) + + // Store directly in thread-blocked order + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ItemsPerThread; ITEM++) + { + thread_itr[ITEM] = items[ITEM]; + } +} + +//! @rst +//! Store a blocked arrangement of items across a +//! thread block into a linear segment of items, guarded by range +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @blocked +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to store. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam OutputIteratorT +//! **[inferred]** The random-access iterator type for output @iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread +//! (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D thread blocks) +//! +//! @param[in] block_itr +//! The thread block's base output iterator for storing to +//! +//! @param[in] items +//! Data to store +//! +//! @param[in] valid_items +//! Number of valid items to write +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +StoreDirectBlocked(int linear_tid, OutputIteratorT block_itr, T (&items)[ItemsPerThread], int valid_items) +{ + OutputIteratorT thread_itr = block_itr + (linear_tid * ItemsPerThread); // NOLINT(bugprone-misplaced-widening-cast) + + // Store directly in thread-blocked order + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ItemsPerThread; ITEM++) + { + if (ITEM + (linear_tid * ItemsPerThread) < valid_items) + { + thread_itr[ITEM] = items[ITEM]; + } + } +} + +//! @rst +//! Store a blocked arrangement of items across a +//! thread block into a linear segment of items. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @blocked +//! +//! The output offset (``block_ptr + block_offset``) must be quad-item aligned, +//! which is the default starting offset returned by ``cudaMalloc()`` +//! +//! The following conditions will prevent vectorization and storing will +//! fall back to cub::BLOCK_STORE_DIRECT: +//! +//! - ``ItemsPerThread`` is odd +//! - The data type ``T`` is not a built-in primitive or CUDA vector type +//! (e.g., ``short``, ``int2``, ``double``, ``float2``, etc.) +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to store. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread +//! (e.g., ``(threadIdx.y * blockDim.x) + linear_tid`` for 2D thread blocks) +//! +//! @param[in] block_ptr +//! Input pointer for storing from +//! +//! @param[in] items +//! Data to store +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +StoreDirectBlockedVectorized(int linear_tid, T* block_ptr, T (&items)[ItemsPerThread]) +{ + // Maximum CUDA vector size is 4 elements + static constexpr int MAX_VEC_SIZE = ::cuda::std::min(4, ItemsPerThread); + + // Vector size must be a power of two and an even divisor of the items per thread + static constexpr int VEC_SIZE = + ((((MAX_VEC_SIZE - 1) & MAX_VEC_SIZE) == 0) && ((ItemsPerThread % MAX_VEC_SIZE) == 0)) ? MAX_VEC_SIZE : 1; + + static constexpr int VECTORS_PER_THREAD = ItemsPerThread / VEC_SIZE; + + // Vector type + using Vector = typename CubVector::Type; + + // Add the alignment check to ensure the vectorized storing can proceed. + if (::cuda::std::is_sufficiently_aligned(block_ptr)) + { + // Alias global pointer + Vector* block_ptr_vectors = reinterpret_cast(const_cast(block_ptr)); + + // Alias pointers (use "raw" array here which should get optimized away to prevent conservative PTXAS lmem spilling) + Vector raw_vector[VECTORS_PER_THREAD]; + T* raw_items = reinterpret_cast(raw_vector); + + // Copy + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ItemsPerThread; ITEM++) + { + raw_items[ITEM] = items[ITEM]; + } + + // Direct-store using vector types + StoreDirectBlocked(linear_tid, block_ptr_vectors, raw_vector); + } + else + { + // Direct-store using original type when the address is misaligned + StoreDirectBlocked(linear_tid, block_ptr, items); + } +} + +//! @} +//! @name Striped arrangement I/O (direct) +//! @{ + +//! @rst +//! Store a striped arrangement of data across the thread block into a +//! linear segment of items. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @striped +//! @endrst +//! +//! @tparam BLOCK_THREADS +//! The thread block size in threads +//! +//! @tparam T +//! **[inferred]** The data type to store. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam OutputIteratorT +//! **[inferred]** The random-access iterator type for output @iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread +//! (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D thread blocks) +//! +//! @param[in] block_itr +//! The thread block's base output iterator for storing to +//! +//! @param[in] items +//! Data to store +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +StoreDirectStriped(int linear_tid, OutputIteratorT block_itr, T (&items)[ItemsPerThread]) +{ + OutputIteratorT thread_itr = block_itr + linear_tid; + + // Store directly in striped order + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ItemsPerThread; ITEM++) + { + thread_itr[(ITEM * BLOCK_THREADS)] = items[ITEM]; + } +} + +//! @rst +//! Store a striped arrangement of data across the thread block into +//! a linear segment of items, guarded by range +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @striped +//! @endrst +//! +//! @tparam BLOCK_THREADS +//! The thread block size in threads +//! +//! @tparam T +//! **[inferred]** The data type to store. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam OutputIteratorT +//! **[inferred]** The random-access iterator type for output @iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread +//! (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D thread blocks) +//! +//! @param[in] block_itr +//! The thread block's base output iterator for storing to +//! +//! @param[in] items +//! Data to store +//! +//! @param[in] valid_items +//! Number of valid items to write +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +StoreDirectStriped(int linear_tid, OutputIteratorT block_itr, T (&items)[ItemsPerThread], int valid_items) +{ + OutputIteratorT thread_itr = block_itr + linear_tid; + + // Store directly in striped order + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ItemsPerThread; ITEM++) + { + if ((ITEM * BLOCK_THREADS) + linear_tid < valid_items) + { + thread_itr[(ITEM * BLOCK_THREADS)] = items[ITEM]; + } + } +} + +//! @} +//! @name Warp-striped arrangement I/O (direct) +//! @{ + +//! @rst +//! Store a warp-striped arrangement of data across the +//! thread block into a linear segment of items. +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @warpstriped +//! +//! Usage Considerations +//! ++++++++++++++++++++ +//! +//! The number of threads in the thread block must be a multiple of the architecture's warp size. +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to store. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam OutputIteratorT +//! **[inferred]** The random-access iterator type for output @iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread +//! (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D thread blocks) +//! +//! @param[in] block_itr +//! The thread block's base output iterator for storing to +//! +//! @param[out] items +//! Data to load +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +StoreDirectWarpStriped(int linear_tid, OutputIteratorT block_itr, T (&items)[ItemsPerThread]) +{ + int tid = linear_tid & (detail::warp_threads - 1); + int wid = linear_tid >> detail::log2_warp_threads; + int warp_offset = wid * detail::warp_threads * ItemsPerThread; + + OutputIteratorT thread_itr = block_itr + warp_offset + tid; + + // Store directly in warp-striped order + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ItemsPerThread; ITEM++) + { + thread_itr[(ITEM * detail::warp_threads)] = items[ITEM]; // NOLINT(bugprone-misplaced-widening-cast) + } +} + +//! @rst +//! Store a warp-striped arrangement of data across the thread block into a +//! linear segment of items, guarded by range +//! +//! .. versionadded:: 2.2.0 +//! First appears in CUDA Toolkit 12.3. +//! +//! @warpstriped +//! +//! Usage Considerations +//! ++++++++++++++++++++ +//! +//! The number of threads in the thread block must be a multiple of the architecture's warp size. +//! @endrst +//! +//! @tparam T +//! **[inferred]** The data type to store. +//! +//! @tparam ItemsPerThread +//! **[inferred]** The number of consecutive items partitioned onto each thread. +//! +//! @tparam OutputIteratorT +//! **[inferred]** The random-access iterator type for output @iterator. +//! +//! @param[in] linear_tid +//! A suitable 1D thread-identifier for the calling thread +//! (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D thread blocks) +//! +//! @param[in] block_itr +//! The thread block's base output iterator for storing to +//! +//! @param[in] items +//! Data to store +//! +//! @param[in] valid_items +//! Number of valid items to write +template +_CCCL_DEVICE _CCCL_FORCEINLINE void +StoreDirectWarpStriped(int linear_tid, OutputIteratorT block_itr, T (&items)[ItemsPerThread], int valid_items) +{ + int tid = linear_tid & (detail::warp_threads - 1); + int wid = linear_tid >> detail::log2_warp_threads; + int warp_offset = wid * detail::warp_threads * ItemsPerThread; + + OutputIteratorT thread_itr = block_itr + warp_offset + tid; + + // Store directly in warp-striped order + _CCCL_PRAGMA_UNROLL_FULL() + for (int ITEM = 0; ITEM < ItemsPerThread; ITEM++) + { + if (warp_offset + tid + (ITEM * detail::warp_threads) < valid_items) + { + thread_itr[(ITEM * detail::warp_threads)] = items[ITEM]; // NOLINT(bugprone-misplaced-widening-cast) + } + } +} + +//! @} + +//----------------------------------------------------------------------------- +// Generic BlockStore abstraction +//----------------------------------------------------------------------------- + +//! cub::BlockStoreAlgorithm enumerates alternative algorithms for cub::BlockStore to write a +//! blocked arrangement of items across a CUDA thread block to a linear segment of memory. +enum BlockStoreAlgorithm +{ + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`blocked arrangement ` of data is written directly to memory. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) decreases as the + //! access stride between threads increases (i.e., the number items per thread). + //! + //! @endrst + BLOCK_STORE_DIRECT, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`striped arrangement ` of data is written directly to memory. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! The utilization of memory transactions (coalescing) remains high regardless + //! of items written per thread. + //! + //! @endrst + BLOCK_STORE_STRIPED, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`blocked arrangement ` of data is written directly + //! to memory using CUDA's built-in vectorized stores as a coalescing optimization. + //! For example, ``st.global.v4.s32`` instructions will be generated + //! when ``T = int`` and ``ItemsPerThread % 4 == 0``. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) remains high until the the + //! access stride between threads (i.e., the number items per thread) exceeds the + //! maximum vector store width (typically 4 items or 64B, whichever is lower). + //! - The following conditions will prevent vectorization and writing will fall back to cub::BLOCK_STORE_DIRECT: + //! + //! - ``ItemsPerThread`` is odd + //! - The ``OutputIteratorT`` is not a simple pointer type + //! - The block output offset is not quadword-aligned + //! - The data type ``T`` is not a built-in primitive or CUDA vector type + //! (e.g., ``short``, ``int2``, ``double``, ``float2``, etc.) + //! + //! @endrst + BLOCK_STORE_VECTORIZE, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`blocked arrangement ` is locally + //! transposed and then efficiently written to memory as a :ref:`striped arrangement `. + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) remains high regardless + //! of items written per thread. + //! - The local reordering incurs slightly longer latencies and throughput than the + //! direct cub::BLOCK_STORE_DIRECT and cub::BLOCK_STORE_VECTORIZE alternatives. + //! + //! @endrst + BLOCK_STORE_TRANSPOSE, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`blocked arrangement ` is locally + //! transposed and then efficiently written to memory as a + //! :ref:`warp-striped arrangement `. + //! + //! Usage Considerations + //! ++++++++++++++++++++++++++ + //! + //! - BLOCK_THREADS must be a multiple of WARP_THREADS + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) remains high regardless + //! of items written per thread. + //! - The local reordering incurs slightly longer latencies and throughput than the + //! direct cub::BLOCK_STORE_DIRECT and cub::BLOCK_STORE_VECTORIZE alternatives. + //! + //! @endrst + BLOCK_STORE_WARP_TRANSPOSE, + + //! @rst + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! A :ref:`blocked arrangement ` is locally + //! transposed and then efficiently written to memory as a + //! :ref:`warp-striped arrangement `. + //! To reduce the shared memory requirement, only one warp's worth of shared + //! memory is provisioned and is subsequently time-sliced among warps. + //! + //! Usage Considerations + //! ++++++++++++++++++++++++++ + //! + //! - BLOCK_THREADS must be a multiple of WARP_THREADS + //! + //! Performance Considerations + //! ++++++++++++++++++++++++++ + //! + //! - The utilization of memory transactions (coalescing) remains high regardless + //! of items written per thread. + //! - Provisions less shared memory temporary storage, but incurs larger + //! latencies than the BLOCK_STORE_WARP_TRANSPOSE alternative. + //! + //! @endrst + BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED, +}; + +#if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED) +namespace detail +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr const char* to_string(BlockStoreAlgorithm algo) noexcept +{ + switch (algo) + { + case BLOCK_STORE_DIRECT: + return "BLOCK_STORE_DIRECT"; + case BLOCK_STORE_STRIPED: + return "BLOCK_STORE_STRIPED"; + case BLOCK_STORE_VECTORIZE: + return "BLOCK_STORE_VECTORIZE"; + case BLOCK_STORE_TRANSPOSE: + return "BLOCK_STORE_TRANSPOSE"; + case BLOCK_STORE_WARP_TRANSPOSE: + return "BLOCK_STORE_WARP_TRANSPOSE"; + case BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED: + return "BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED"; + } + return ""; +} +} // namespace detail + +inline ::std::ostream& operator<<(::std::ostream& os, BlockStoreAlgorithm algo) +{ + return os << CUB_NS_QUALIFIER::detail::to_string(algo); +} +#endif // _CCCL_HOSTED() && !_CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END + +#if __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) +template <::cuda::std::same_as CharT> +struct std::formatter : formatter +{ + template + auto format(const CUB_NS_QUALIFIER::BlockStoreAlgorithm& algo, FmtCtx& ctx) const + { + return formatter::format(CUB_NS_QUALIFIER::detail::to_string(algo), ctx); + } +}; +#endif // __cpp_lib_format >= 201907L && !defined(_CCCL_DOXYGEN_INVOKED) + +CUB_NAMESPACE_BEGIN + +//! @rst +//! The BlockStore class provides :ref:`collective ` data movement +//! methods for writing a :ref:`blocked arrangement ` of items +//! partitioned across a CUDA thread block to a linear segment of memory. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! - The BlockStore class provides a single data movement abstraction that can be specialized +//! to implement different cub::BlockStoreAlgorithm strategies. This facilitates different +//! performance policies for different architectures, data types, granularity sizes, etc. +//! - BlockStore can be optionally specialized by different data movement strategies: +//! +//! #. :cpp:enumerator:`cub::BLOCK_STORE_DIRECT`: +//! A :ref:`blocked arrangement ` of data is written directly to memory. +//! #. :cpp:enumerator:`cub::BLOCK_STORE_STRIPED`: +//! A :ref:`striped arrangement ` of data is written directly to memory. +//! #. :cpp:enumerator:`cub::BLOCK_STORE_VECTORIZE`: +//! A :ref:`blocked arrangement ` of data is written directly to memory +//! using CUDA's built-in vectorized stores as a coalescing optimization. +//! #. :cpp:enumerator:`cub::BLOCK_STORE_TRANSPOSE`: +//! A :ref:`blocked arrangement ` is locally transposed into +//! a :ref:`striped arrangement ` which is then written to memory. +//! #. :cpp:enumerator:`cub::BLOCK_STORE_WARP_TRANSPOSE`: +//! A :ref:`blocked arrangement ` is locally transposed into +//! a :ref:`warp-striped arrangement ` which is then written to memory. +//! #. :cpp:enumerator:`cub::BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED`: +//! A :ref:`blocked arrangement ` is locally transposed into +//! a :ref:`warp-striped arrangement ` which is then written to memory. +//! To reduce the shared memory requireent, only one warp's worth of shared memory is provisioned and is +//! subsequently time-sliced among warps. +//! +//! - @rowmajor +//! +//! A Simple Example +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @blockcollective{BlockStore} +//! +//! The code snippet below illustrates the storing of a "blocked" arrangement +//! of 512 integers across 128 threads (where each thread owns 4 consecutive items) +//! into a linear segment of memory. The store is specialized for ``BLOCK_STORE_WARP_TRANSPOSE``, +//! meaning items are locally reordered among threads so that memory references will be +//! efficiently coalesced using a warp-striped access pattern. +//! +//! .. code-block:: c++ +//! +//! #include // or equivalently +//! +//! __global__ void ExampleKernel(int *d_data, ...) +//! { +//! // Specialize BlockStore for a 1D block of 128 threads owning 4 integer items each +//! using BlockStore = cub::BlockStore; +//! +//! // Allocate shared memory for BlockStore +//! __shared__ typename BlockStore::TempStorage temp_storage; +//! +//! // Obtain a segment of consecutive items that are blocked across threads +//! int thread_data[4]; +//! ... +//! +//! // Store items to linear memory +//! BlockStore(temp_storage).Store(d_data, thread_data); +//! } +//! +//! Suppose the set of ``thread_data`` across the block of threads is +//! ``{ [0,1,2,3], [4,5,6,7], ..., [508,509,510,511] }``. +//! The output ``d_data`` will be ``0, 1, 2, 3, 4, 5, ...``. +//! +//! Re-using dynamically allocating shared memory +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``block/example_block_reduce_dyn_smem.cu`` example illustrates usage of +//! dynamically shared memory with BlockReduce and how to re-purpose the same memory region. +//! This example can be easily adapted to the storage required by BlockStore. +//! +//! @endrst +//! +//! @tparam T +//! The type of data to be written. +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam ItemsPerThread +//! The number of consecutive items partitioned onto each thread. +//! +//! @tparam Algorithm +//! **[optional]** cub::BlockStoreAlgorithm tuning policy enumeration (default: cub::BLOCK_STORE_DIRECT) +//! +//! @tparam BlockDimY +//! **[optional]** The thread block length in threads along the Y dimension (default: 1) +//! +//! @tparam BlockDimZ +//! **[optional]** The thread block length in threads along the Z dimension (default: 1) +//! +template +class BlockStore +{ + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + // transposing store algorithms need a BlockExchange + using block_exchange = + BlockExchange; + + static_assert((Algorithm != BLOCK_STORE_WARP_TRANSPOSE && Algorithm != BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED) + || (BLOCK_THREADS % detail::warp_threads == 0), + "Threads per block must be a multiple of warp_threads for this BlockStoreAlgorithm"); + + _CCCL_HOST_DEVICE_API static constexpr auto temp_storage_helper() + { + if constexpr (Algorithm == BLOCK_STORE_DIRECT || Algorithm == BLOCK_STORE_STRIPED + || Algorithm == BLOCK_STORE_VECTORIZE) + { + return NullType{}; + } + else if constexpr (Algorithm == BLOCK_STORE_TRANSPOSE || Algorithm == BLOCK_STORE_WARP_TRANSPOSE + || Algorithm == BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED) + { + struct _TempStorage : block_exchange::TempStorage + { + volatile int valid_items; // Temporary storage for partially-full block guard + }; + return _TempStorage{}; + } + } + + using _TempStorage = decltype(temp_storage_helper()); + + _CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage() + { + __shared__ _TempStorage private_storage; + return private_storage; + } + + _TempStorage& temp_storage; + + int linear_tid; + +public: + //! @smemstorage{BlockStore} + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //! @name Collective constructors + //! @{ + + /** + * @brief Collective constructor using a private static allocation of shared memory as temporary storage. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockStore() + : temp_storage(PrivateStorage()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @brief Collective constructor using the specified memory allocation as temporary storage. + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + * + * @param[in] temp_storage + * Reference to memory allocation having layout type TempStorage + */ + _CCCL_DEVICE _CCCL_FORCEINLINE BlockStore(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //! @} + //! @name Data movement + //! @{ + + //! @rst + //! Store items into a linear segment of memory + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @blocked + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the storing of a "blocked" arrangement + //! of 512 integers across 128 threads (where each thread owns 4 consecutive items) + //! into a linear segment of memory. The store is specialized for ``BLOCK_STORE_WARP_TRANSPOSE``, + //! meaning items are locally reordered among threads so that memory references will be + //! efficiently coalesced using a warp-striped access pattern. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, ...) + //! { + //! // Specialize BlockStore for a 1D block of 128 threads owning 4 integer items each + //! using BlockStore = cub::BlockStore; + //! + //! // Allocate shared memory for BlockStore + //! __shared__ typename BlockStore::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Store items to linear memory + //! BlockStore(temp_storage).Store(d_data, thread_data); + //! } + //! + //! Suppose the set of ``thread_data`` across the block of threads is + //! ``{ [0,1,2,3], [4,5,6,7], ..., [508,509,510,511] }``. + //! The output ``d_data`` will be ``0, 1, 2, 3, 4, 5, ...``. + //! @endrst + //! + //! @param[out] block_itr + //! The thread block's base output iterator for storing to + //! + //! @param[in] items + //! Data to store + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Store(OutputIteratorT block_itr, T (&items)[ItemsPerThread]) + { + if constexpr (Algorithm == BLOCK_STORE_DIRECT) + { + StoreDirectBlocked(linear_tid, block_itr, items); + } + else if constexpr (Algorithm == BLOCK_STORE_STRIPED) + { + StoreDirectStriped(linear_tid, block_itr, items); + } + else if constexpr (Algorithm == BLOCK_STORE_VECTORIZE) + { + if constexpr (::cuda::std::contiguous_iterator && ::cuda::std::__can_to_address) + { + StoreDirectBlockedVectorized(linear_tid, ::cuda::std::to_address(block_itr), items); + } + else + { + StoreDirectBlocked(linear_tid, block_itr, items); + } + } + else if constexpr (Algorithm == BLOCK_STORE_TRANSPOSE) + { + block_exchange(temp_storage).BlockedToStriped(items); + StoreDirectStriped(linear_tid, block_itr, items); + } + else if constexpr (Algorithm == BLOCK_STORE_WARP_TRANSPOSE || Algorithm == BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED) + { + block_exchange(temp_storage).BlockedToWarpStriped(items); + StoreDirectWarpStriped(linear_tid, block_itr, items); + } + } + + //! @rst + //! Store items into a linear segment of memory, guarded by range. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - @blocked + //! - @smemreuse + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the guarded storing of a "blocked" arrangement + //! of 512 integers across 128 threads (where each thread owns 4 consecutive items) + //! into a linear segment of memory. The store is specialized for ``BLOCK_STORE_WARP_TRANSPOSE``, + //! meaning items are locally reordered among threads so that memory references will be + //! efficiently coalesced using a warp-striped access pattern. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! __global__ void ExampleKernel(int *d_data, int valid_items, ...) + //! { + //! // Specialize BlockStore for a 1D block of 128 threads owning 4 integer items each + //! using BlockStore = cub::BlockStore; + //! + //! // Allocate shared memory for BlockStore + //! __shared__ typename BlockStore::TempStorage temp_storage; + //! + //! // Obtain a segment of consecutive items that are blocked across threads + //! int thread_data[4]; + //! ... + //! + //! // Store items to linear memory + //! BlockStore(temp_storage).Store(d_data, thread_data, valid_items); + //! } + //! + //! Suppose the set of ``thread_data`` across the block of threads is + //! ``{ [0,1,2,3], [4,5,6,7], ..., [508,509,510,511] }`` and ``valid_items`` is ``5``. + //! The output ``d_data`` will be ``0, 1, 2, 3, 4, ?, ?, ?, ...``, with + //! only the first two threads being unmasked to store portions of valid data. + //! @endrst + //! + //! @param[out] block_itr + //! The thread block's base output iterator for storing to + //! + //! @param[in] items + //! Data to store + //! + //! @param[in] valid_items + //! Number of valid items to write + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Store(OutputIteratorT block_itr, T (&items)[ItemsPerThread], int valid_items) + { + if constexpr (Algorithm == BLOCK_STORE_DIRECT || Algorithm == BLOCK_STORE_VECTORIZE) + { + StoreDirectBlocked(linear_tid, block_itr, items, valid_items); + } + else if constexpr (Algorithm == BLOCK_STORE_STRIPED) + { + StoreDirectStriped(linear_tid, block_itr, items, valid_items); + } + else if constexpr (Algorithm == BLOCK_STORE_TRANSPOSE) + { + block_exchange(temp_storage).BlockedToStriped(items); + if (linear_tid == 0) + { + // Move through volatile smem as a workaround to prevent RF spilling on subsequent loads + temp_storage.valid_items = valid_items; + } + __syncthreads(); + StoreDirectStriped(linear_tid, block_itr, items, temp_storage.valid_items); + } + else if constexpr (Algorithm == BLOCK_STORE_WARP_TRANSPOSE || Algorithm == BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED) + { + block_exchange(temp_storage).BlockedToWarpStriped(items); + if (linear_tid == 0) + { + // Move through volatile smem as a workaround to prevent RF spilling on subsequent loads + temp_storage.valid_items = valid_items; + } + __syncthreads(); + StoreDirectWarpStriped(linear_tid, block_itr, items, temp_storage.valid_items); + } + } + + //! @} +}; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document +template > +struct BlockStoreType +{ + using type = cub::BlockStore; +}; +#endif // _CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/block_topk.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/block_topk.cuh new file mode 100644 index 00000000..1a56e655 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/block_topk.cuh @@ -0,0 +1,89 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +// TODO (elstehle): Add documentation +template +class block_topk +{ +private: + using internal_block_topk_t = block_topk_air; + +public: + struct TempStorage + { + typename internal_block_topk_t::TempStorage topk_storage; + }; + +private: + TempStorage& storage; + +public: + _CCCL_DEVICE_API _CCCL_FORCEINLINE block_topk(TempStorage& storage) + : storage(storage) + {} + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void max_pairs( + KeyT (&keys)[ItemsPerThread], + ValueT (&values)[ItemsPerThread], + int k, + int num_valid, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8) + { + internal_block_topk_t(storage.topk_storage) + .template select_pairs(keys, values, k, num_valid, begin_bit, end_bit); + } + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void + max_keys(KeyT (&keys)[ItemsPerThread], int k, int num_valid, int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + internal_block_topk_t(storage.topk_storage) + .template select_keys(keys, k, num_valid, begin_bit, end_bit); + } + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void min_pairs( + KeyT (&keys)[ItemsPerThread], + ValueT (&values)[ItemsPerThread], + int k, + int num_valid, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8) + { + internal_block_topk_t(storage.topk_storage) + .template select_pairs(keys, values, k, num_valid, begin_bit, end_bit); + } + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void + min_keys(KeyT (&keys)[ItemsPerThread], int k, int num_valid, int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + internal_block_topk_t(storage.topk_storage) + .template select_keys(keys, k, num_valid, begin_bit, end_bit); + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/radix_rank_sort_operations.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/radix_rank_sort_operations.cuh new file mode 100644 index 00000000..6a47d209 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/radix_rank_sort_operations.cuh @@ -0,0 +1,580 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2020, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * radix_rank_sort_operations.cuh contains common abstractions, definitions and + * operations used for radix sorting and ranking. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +/** \brief Base struct for digit extractor. Contains common code to provide + special handling for floating-point -0.0. + + \note This handles correctly both the case when the keys are + bitwise-complemented after twiddling for descending sort (in onesweep) as + well as when the keys are not bit-negated, but the implementation handles + descending sort separately (in other implementations in CUB). Twiddling + alone maps -0.0f to 0x7fffffff and +0.0f to 0x80000000 for float, which are + subsequent bit patterns and bitwise complements of each other. For onesweep, + both -0.0f and +0.0f are mapped to the bit pattern of +0.0f (0x80000000) for + ascending sort, and to the pattern of -0.0f (0x7fffffff) for descending + sort. For all other sorting implementations in CUB, both are always mapped + to +0.0f. Since bit patterns for both -0.0f and +0.0f are next to each other + and only one of them is used, the sorting works correctly. For double, the + same applies, but with 64-bit patterns. +*/ +template > +struct BaseDigitExtractor +{ + using TraitsT = Traits; + using UnsignedBits = typename TraitsT::UnsignedBits; + + static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE UnsignedBits ProcessFloatMinusZero(UnsignedBits key) + { + return key; + } +}; + +template +struct BaseDigitExtractor +{ + using TraitsT = Traits; + using UnsignedBits = typename TraitsT::UnsignedBits; + + static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE UnsignedBits ProcessFloatMinusZero(UnsignedBits key) + { + UnsignedBits TWIDDLED_MINUS_ZERO_BITS = + TraitsT::TwiddleIn(UnsignedBits(1) << UnsignedBits(8 * sizeof(UnsignedBits) - 1)); + UnsignedBits TWIDDLED_ZERO_BITS = TraitsT::TwiddleIn(0); + return key == TWIDDLED_MINUS_ZERO_BITS ? TWIDDLED_ZERO_BITS : key; + } +}; + +/** \brief A wrapper type to extract digits. Uses the BFE intrinsic to extract a + * key from a digit. */ +template +struct BFEDigitExtractor : BaseDigitExtractor +{ + using typename BaseDigitExtractor::UnsignedBits; + + ::cuda::std::uint32_t bit_start; + ::cuda::std::uint32_t num_bits; + + explicit _CCCL_DEVICE _CCCL_FORCEINLINE + BFEDigitExtractor(::cuda::std::uint32_t bit_start = 0, ::cuda::std::uint32_t num_bits = 0) + : bit_start(bit_start) + , num_bits(num_bits) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t Digit(UnsignedBits key) const + { + return ::cuda::bitfield_extract(this->ProcessFloatMinusZero(key), bit_start, num_bits); + } +}; + +/** \brief A wrapper type to extract digits. Uses a combination of shift and + * bitwise and to extract digits. */ +template +struct ShiftDigitExtractor : BaseDigitExtractor +{ + using typename BaseDigitExtractor::UnsignedBits; + + ::cuda::std::uint32_t bit_start; + ::cuda::std::uint32_t mask; // NOLINT(modernize-use-default-member-init) + + explicit _CCCL_DEVICE _CCCL_FORCEINLINE + ShiftDigitExtractor(::cuda::std::uint32_t bit_start = 0, ::cuda::std::uint32_t num_bits = 0) + : bit_start(bit_start) + , mask((1 << num_bits) - 1) + {} + + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t Digit(UnsignedBits key) const + { + return ::cuda::std::uint32_t(this->ProcessFloatMinusZero(key) >> UnsignedBits(bit_start)) & mask; + } +}; + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document +namespace detail +{ +struct identity_decomposer_t +{ + template + _CCCL_HOST_DEVICE T& operator()(T& key) const + { + return key; + } +}; + +template +_CCCL_HOST_DEVICE void +for_each_member_impl(F f, const ::cuda::std::tuple& tpl, ::cuda::std::index_sequence) +{ + static_assert(sizeof...(Ts), "Empty aggregates are not supported"); + + // Most radix operations are indifferent to the order of operations. Conversely, the digit extractor traverses fields + // from the least significant to the most significant to imitate bitset printing where higher bits are on the left. It + // also maps to intuition, where something coming first is more important. Therefore, we traverse fields on the + // opposite order. + + // we use a fold over the assignment operator to get right-to-left evaluation order + [[maybe_unused]] int dummy; + ((f(::cuda::std::get(tpl)), dummy) = ... = 0); +} + +template +_CCCL_HOST_DEVICE void for_each_member(F f, DecomposerT decomposer, T& aggregate) +{ + const auto& tuple_of_refs = decomposer(aggregate); + constexpr int tuple_size = ::cuda::std::tuple_size_v<::cuda::std::remove_reference_t>; + for_each_member_impl(f, tuple_of_refs, ::cuda::std::make_index_sequence{}); +} + +namespace radix +{ +// True for types that can be converted to bit ordered values using cub::Traits::UnsignedBits (and TwiddleIn/Out) +template +inline constexpr bool can_twiddle = false; + +template +inline constexpr bool can_twiddle::UnsignedBits>> = true; + +template +inline constexpr bool can_twiddle_tuple_refs = false; + +template +inline constexpr bool can_twiddle_tuple_refs<::cuda::std::tuple> = (can_twiddle && ...); + +template +inline constexpr bool decomposer_check = can_twiddle_tuple_refs<::cuda::std::invoke_result_t>; + +// SFINAE-friendly version of decomposer_check_t: true iff DecomposerT is callable +// with KeyT& and returns a tuple of references to fundamental types. +template +inline constexpr bool is_valid_decomposer = false; + +template +inline constexpr bool + is_valid_decomposer>> = + can_twiddle_tuple_refs<::cuda::std::invoke_result_t>; + +template +struct bit_ordered_conversion_policy_t +{ + using bit_ordered_type = typename Traits::UnsignedBits; + + static _CCCL_HOST_DEVICE bit_ordered_type to_bit_ordered(detail::identity_decomposer_t, bit_ordered_type val) + { + return Traits::TwiddleIn(val); + } + + static _CCCL_HOST_DEVICE bit_ordered_type from_bit_ordered(detail::identity_decomposer_t, bit_ordered_type val) + { + return Traits::TwiddleOut(val); + } +}; + +template +struct bit_ordered_inversion_policy_t +{ + using bit_ordered_type = typename Traits::UnsignedBits; + + static _CCCL_HOST_DEVICE bit_ordered_type inverse(detail::identity_decomposer_t, bit_ordered_type val) + { + return ~val; + } +}; + +template > +struct traits_t +{ + using bit_ordered_type = typename Traits::UnsignedBits; + using bit_ordered_conversion_policy = bit_ordered_conversion_policy_t; + using bit_ordered_inversion_policy = bit_ordered_inversion_policy_t; + + template + using digit_extractor_t = FundamentalExtractorT; + + static _CCCL_HOST_DEVICE bit_ordered_type min_raw_binary_key(detail::identity_decomposer_t) + { + return Traits::LOWEST_KEY; + } + + static _CCCL_HOST_DEVICE bit_ordered_type max_raw_binary_key(detail::identity_decomposer_t) + { + return Traits::MAX_KEY; + } + + static _CCCL_HOST_DEVICE int default_end_bit(detail::identity_decomposer_t) + { + return sizeof(T) * 8; + } + + template + static _CCCL_HOST_DEVICE digit_extractor_t + digit_extractor(int begin_bit, int num_bits, detail::identity_decomposer_t) + { + return FundamentalExtractorT(begin_bit, num_bits); + } +}; + +template +struct traits_t : traits_t +{}; + +template +struct min_raw_binary_key_f +{ + DecomposerT decomposer; + + template + _CCCL_HOST_DEVICE void operator()(T& field) + { + using traits = traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + // TODO(bgruber): was it intended to pass decomposer here instead of identity_decomposer_t? + reinterpret_cast(field) = traits::min_raw_binary_key(detail::identity_decomposer_t{}); + } +}; + +template +_CCCL_HOST_DEVICE void min_raw_binary_key(DecomposerT decomposer, T& aggregate) +{ + detail::for_each_member(min_raw_binary_key_f{decomposer}, decomposer, aggregate); +} + +template +struct max_raw_binary_key_f +{ + DecomposerT decomposer; + + template + _CCCL_HOST_DEVICE void operator()(T& field) + { + using traits = traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + // TODO(bgruber): was it intended to pass decomposer here instead of identity_decomposer_t? + reinterpret_cast(field) = traits::max_raw_binary_key(detail::identity_decomposer_t{}); + } +}; + +template +_CCCL_HOST_DEVICE void max_raw_binary_key(DecomposerT decomposer, T& aggregate) +{ + detail::for_each_member(max_raw_binary_key_f{decomposer}, decomposer, aggregate); +} + +template +struct to_bit_ordered_f +{ + DecomposerT decomposer; + + template + _CCCL_HOST_DEVICE void operator()(T& field) + { + using traits = traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion = typename traits::bit_ordered_conversion_policy; + + auto& ordered_field = reinterpret_cast(field); + // TODO(bgruber): was it intended to pass decomposer here instead of identity_decomposer_t? + ordered_field = bit_ordered_conversion::to_bit_ordered(detail::identity_decomposer_t{}, ordered_field); + } +}; + +template +_CCCL_HOST_DEVICE void to_bit_ordered(DecomposerT decomposer, T& aggregate) +{ + detail::for_each_member(to_bit_ordered_f{decomposer}, decomposer, aggregate); +} + +template +struct from_bit_ordered_f +{ + DecomposerT decomposer; + + template + _CCCL_HOST_DEVICE void operator()(T& field) + { + using traits = traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion = typename traits::bit_ordered_conversion_policy; + + auto& ordered_field = reinterpret_cast(field); + // TODO(bgruber): was it intended to pass decomposer here instead of identity_decomposer_t? + ordered_field = bit_ordered_conversion::from_bit_ordered(detail::identity_decomposer_t{}, ordered_field); + } +}; + +template +_CCCL_HOST_DEVICE void from_bit_ordered(DecomposerT decomposer, T& aggregate) +{ + detail::for_each_member(from_bit_ordered_f{decomposer}, decomposer, aggregate); +} + +struct inverse_f +{ + template + _CCCL_HOST_DEVICE void operator()(T& field) + { + using traits = traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + + auto& ordered_field = reinterpret_cast(field); + ordered_field = ~ordered_field; + } +}; + +template +_CCCL_HOST_DEVICE void inverse(DecomposerT decomposer, T& aggregate) +{ + detail::for_each_member(inverse_f{}, decomposer, aggregate); +} + +struct default_end_bit_f +{ + int& result; + + template + _CCCL_HOST_DEVICE void operator()(T& /* field */) + { + result += sizeof(T) * 8; + } +}; + +template +_CCCL_HOST_DEVICE int default_end_bit(DecomposerT decomposer, T& aggregate) +{ + int result{}; + detail::for_each_member(default_end_bit_f{result}, decomposer, aggregate); + return result; +} + +struct digit_f +{ + ::cuda::std::uint32_t& dst; + ::cuda::std::uint32_t& dst_bit_start; + ::cuda::std::uint32_t& src_bit_start; + ::cuda::std::uint32_t& num_bits; + + template + _CCCL_HOST_DEVICE void operator()(T& src) + { + constexpr ::cuda::std::uint32_t src_size = sizeof(T) * 8; + + if (src_bit_start >= src_size) + { + src_bit_start -= src_size; + } + else + { + using traits = traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + + const ::cuda::std::uint32_t bits_to_copy = (::cuda::std::min) (src_size - src_bit_start, num_bits); + + if (bits_to_copy) + { + bit_ordered_type ordered_src = + BaseDigitExtractor::ProcessFloatMinusZero(reinterpret_cast(src)); + + const ::cuda::std::uint32_t mask = (1 << bits_to_copy) - 1; + dst = dst | (((ordered_src >> src_bit_start) & mask) << dst_bit_start); + + num_bits -= bits_to_copy; + dst_bit_start += bits_to_copy; + } + src_bit_start = 0; + } + } +}; + +template +_CCCL_HOST_DEVICE void +digit(DecomposerT decomposer, + ::cuda::std::uint32_t& dst, + T& src, + ::cuda::std::uint32_t& dst_bit_start, + ::cuda::std::uint32_t& src_bit_start, + ::cuda::std::uint32_t& num_bits) +{ + detail::for_each_member(digit_f{dst, dst_bit_start, src_bit_start, num_bits}, decomposer, src); +} + +template +struct custom_digit_extractor_t +{ + DecomposerT decomposer; + ::cuda::std::uint32_t bit_start; + ::cuda::std::uint32_t num_bits; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE + custom_digit_extractor_t(DecomposerT decomposer, ::cuda::std::uint32_t bit_start, ::cuda::std::uint32_t num_bits) + : decomposer(decomposer) + , bit_start(bit_start) + , num_bits(num_bits) + {} + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE ::cuda::std::uint32_t Digit(T& key) const + { + ::cuda::std::uint32_t result{}; + ::cuda::std::uint32_t dst_bit_start{}; + ::cuda::std::uint32_t src_bit_start = bit_start; + ::cuda::std::uint32_t bits_remaining{num_bits}; + digit(decomposer, result, key, dst_bit_start, src_bit_start, bits_remaining); + return result; + } +}; + +struct custom_bit_conversion_policy_t +{ + template + static _CCCL_HOST_DEVICE T to_bit_ordered(DecomposerT decomposer, T val) + { + detail::radix::to_bit_ordered(decomposer, val); + return val; + } + + template + static _CCCL_HOST_DEVICE T from_bit_ordered(DecomposerT decomposer, T val) + { + detail::radix::from_bit_ordered(decomposer, val); + return val; + } +}; + +struct custom_bit_inversion_policy_t +{ + template + static _CCCL_HOST_DEVICE T inverse(DecomposerT decomposer, T val) + { + detail::radix::inverse(decomposer, val); + return val; + } +}; + +template +struct traits_t +{ + using bit_ordered_type = T; + using bit_ordered_conversion_policy = custom_bit_conversion_policy_t; + using bit_ordered_inversion_policy = custom_bit_inversion_policy_t; + + template + using digit_extractor_t = custom_digit_extractor_t; + + template + static _CCCL_HOST_DEVICE bit_ordered_type min_raw_binary_key(DecomposerT decomposer) + { + T val{}; + detail::radix::min_raw_binary_key(decomposer, val); + return val; + } + + template + static _CCCL_HOST_DEVICE bit_ordered_type max_raw_binary_key(DecomposerT decomposer) + { + T val{}; + detail::radix::max_raw_binary_key(decomposer, val); + return val; + } + + template + static _CCCL_HOST_DEVICE int default_end_bit(DecomposerT decomposer) + { + T aggregate{}; + return detail::radix::default_end_bit(decomposer, aggregate); + } + + template + static _CCCL_HOST_DEVICE digit_extractor_t + digit_extractor(int begin_bit, int num_bits, DecomposerT decomposer) + { + return custom_digit_extractor_t(decomposer, begin_bit, num_bits); + } +}; +} // namespace radix +} // namespace detail +#endif // _CCCL_DOXYGEN_INVOKED + +//! Twiddling keys for radix sort +template +struct RadixSortTwiddle +{ +private: + using traits = detail::radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion_policy = typename traits::bit_ordered_conversion_policy; + using bit_ordered_inversion_policy = typename traits::bit_ordered_inversion_policy; + +public: + template + static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE // + bit_ordered_type + In(bit_ordered_type key, DecomposerT decomposer = {}) + { + key = bit_ordered_conversion_policy::to_bit_ordered(decomposer, key); + if constexpr (IS_DESCENDING) + { + key = bit_ordered_inversion_policy::inverse(decomposer, key); + } + return key; + } + + template + static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE // + bit_ordered_type + Out(bit_ordered_type key, DecomposerT decomposer = {}) + { + if constexpr (IS_DESCENDING) + { + key = bit_ordered_inversion_policy::inverse(decomposer, key); + } + key = bit_ordered_conversion_policy::from_bit_ordered(decomposer, key); + return key; + } + + template + static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE // + bit_ordered_type + DefaultKey(DecomposerT decomposer = {}) + { + return IS_DESCENDING ? traits::min_raw_binary_key(decomposer) : traits::max_raw_binary_key(decomposer); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_histogram_atomic.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_histogram_atomic.cuh new file mode 100644 index 00000000..102d72df --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_histogram_atomic.cuh @@ -0,0 +1,62 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * The cub::BlockHistogramAtomic class provides atomic-based methods for constructing block-wide + * histograms from data samples partitioned across a CUDA thread block. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +CUB_NAMESPACE_BEGIN +namespace detail +{ +/** + * @brief The BlockHistogramAtomic class provides atomic-based methods for constructing block-wide + * histograms from data samples partitioned across a CUDA thread block. + */ +template +struct BlockHistogramAtomic +{ + /// Shared memory storage layout type + struct TempStorage + {}; + + /// Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE BlockHistogramAtomic(TempStorage& temp_storage) {} + + /** + * @brief Composite data onto an existing histogram + * + * @param[in] items + * Calling thread's input values to histogram + * + * @param[out] histogram + * Reference to shared/device-accessible memory histogram + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Composite(T (&items)[ITEMS_PER_THREAD], CounterT histogram[Bins]) + { + // Update histogram + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < ITEMS_PER_THREAD; ++i) + { + atomicAdd_block(histogram + items[i], 1); + } + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_histogram_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_histogram_sort.cuh new file mode 100644 index 00000000..f0a87b1f --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_histogram_sort.cuh @@ -0,0 +1,209 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * The cub::BlockHistogramSort class provides sorting-based methods for constructing block-wide + * histograms from data samples partitioned across a CUDA thread block. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +/** + * @brief The BlockHistogramSort class provides sorting-based methods for constructing block-wide + * histograms from data samples partitioned across a CUDA thread block. + * + * @tparam T + * Sample type + * + * @tparam BlockDimX + * The thread block length in threads along the X dimension + * + * @tparam ItemsPerThread + * The number of samples per thread + * + * @tparam Bins + * The number of bins into which histogram samples may fall + * + * @tparam BlockDimY + * The thread block length in threads along the Y dimension + * + * @tparam BlockDimZ + * The thread block length in threads along the Z dimension + */ +template +struct BlockHistogramSort +{ + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + // Parameterize BlockRadixSort type for our thread block + using BlockRadixSortT = + BlockRadixSort; + + // Parameterize BlockDiscontinuity type for our thread block + using BlockDiscontinuityT = BlockDiscontinuity; + + /// Shared memory + union _TempStorage + { + // Storage for sorting bin values + typename BlockRadixSortT::TempStorage sort; + + struct Discontinuities + { + // Storage for detecting discontinuities in the tile of sorted bin values + typename BlockDiscontinuityT::TempStorage flag; + + // Storage for noting begin/end offsets of bin runs in the tile of sorted bin values + unsigned int run_begin[Bins]; + unsigned int run_end[Bins]; + } discontinuities; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + // Thread fields + _TempStorage& temp_storage; + unsigned int linear_tid; + + /// Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE BlockHistogramSort(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + // Discontinuity functor + struct DiscontinuityOp + { + // Reference to temp_storage + _TempStorage& temp_storage; + + // Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE DiscontinuityOp(_TempStorage& temp_storage) + : temp_storage(temp_storage) + {} + + // Discontinuity predicate + _CCCL_DEVICE _CCCL_FORCEINLINE bool operator()(const T& a, const T& b, int b_index) + { + if (a != b) + { + // Note the begin/end offsets in shared storage + temp_storage.discontinuities.run_begin[b] = b_index; + temp_storage.discontinuities.run_end[a] = b_index; + + return true; + } + else + { + return false; + } + } + }; + + /** + * @brief Composite data onto an existing histogram + * + * @param[in] items + * Calling thread's input values to histogram + * + * @param[out] histogram + * Reference to shared/device-accessible memory histogram + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void Composite(T (&items)[ItemsPerThread], CounterT histogram[Bins]) + { + static constexpr int TILE_SIZE = BLOCK_THREADS * ItemsPerThread; + + // Sort bytes in blocked arrangement + BlockRadixSortT(temp_storage.sort).Sort(items); + + __syncthreads(); + + // Initialize the shared memory's run_begin and run_end for each bin + int histo_offset = 0; + + _CCCL_PRAGMA_UNROLL_FULL() + for (; histo_offset + BLOCK_THREADS <= Bins; histo_offset += BLOCK_THREADS) + { + temp_storage.discontinuities.run_begin[histo_offset + linear_tid] = TILE_SIZE; + temp_storage.discontinuities.run_end[histo_offset + linear_tid] = TILE_SIZE; + } + // Finish up with guarded initialization if necessary + if ((Bins % BLOCK_THREADS != 0) && (histo_offset + linear_tid < Bins)) + { + temp_storage.discontinuities.run_begin[histo_offset + linear_tid] = TILE_SIZE; + temp_storage.discontinuities.run_end[histo_offset + linear_tid] = TILE_SIZE; + } + + __syncthreads(); + + int flags[ItemsPerThread]; // unused + + // Compute head flags to demarcate contiguous runs of the same bin in the sorted tile + DiscontinuityOp flag_op(temp_storage); + BlockDiscontinuityT(temp_storage.discontinuities.flag).FlagHeads(flags, items, flag_op); + + // Update begin for first item + if (linear_tid == 0) + { + temp_storage.discontinuities.run_begin[items[0]] = 0; + } + + __syncthreads(); + + // Composite into histogram + histo_offset = 0; + + _CCCL_PRAGMA_UNROLL_FULL() + for (; histo_offset + BLOCK_THREADS <= Bins; histo_offset += BLOCK_THREADS) + { + int thread_offset = histo_offset + linear_tid; + CounterT count = + temp_storage.discontinuities.run_end[thread_offset] - temp_storage.discontinuities.run_begin[thread_offset]; + histogram[thread_offset] += count; + } + + // Finish up with guarded composition if necessary + if ((Bins % BLOCK_THREADS != 0) && (histo_offset + linear_tid < Bins)) + { + int thread_offset = histo_offset + linear_tid; + CounterT count = + temp_storage.discontinuities.run_end[thread_offset] - temp_storage.discontinuities.run_begin[thread_offset]; + histogram[thread_offset] += count; + } + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_raking.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_raking.cuh new file mode 100644 index 00000000..2b740da8 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_raking.cuh @@ -0,0 +1,230 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::BlockReduceRaking provides raking-based methods of parallel reduction across a CUDA thread + * block. Supports non-commutative reduction operators. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +/** + * @brief BlockReduceRaking provides raking-based methods of parallel reduction across a CUDA thread + * block. Supports non-commutative reduction operators. + * + * Supports non-commutative binary reduction operators. Unlike commutative + * reduction operators (e.g., addition), the application of a non-commutative + * reduction operator (e.g, string concatenation) across a sequence of inputs must + * honor the relative ordering of items and partial reductions when applying the + * reduction operator. + * + * Compared to the implementation of BlockReduceRakingCommutativeOnly (which + * does not support non-commutative operators), this implementation requires a + * few extra rounds of inter-thread communication. + * + * @tparam T + * Data type being reduced + * + * @tparam BlockDimX + * The thread block length in threads along the X dimension + * + * @tparam BlockDimY + * The thread block length in threads along the Y dimension + * + * @tparam BlockDimZ + * The thread block length in threads along the Z dimension + */ +template +struct BlockReduceRaking +{ + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + /// Layout type for padded thread block raking grid + using BlockRakingLayout = BlockRakingLayout; + + /// WarpReduce utility type + using WarpReduce = typename WarpReduce::InternalWarpReduce; + + /// Constants + /// Number of raking threads + static constexpr int RAKING_THREADS = BlockRakingLayout::RAKING_THREADS; + + /// Number of raking elements per warp synchronous raking thread + static constexpr int SEGMENT_LENGTH = BlockRakingLayout::SEGMENT_LENGTH; + + /// Cooperative work can be entirely warp synchronous + static constexpr bool WARP_SYNCHRONOUS = (RAKING_THREADS == BLOCK_THREADS); + + /// Whether or not warp-synchronous reduction should be unguarded (i.e., the warp-reduction elements is a power of + /// two + static constexpr int WARP_SYNCHRONOUS_UNGUARDED = ::cuda::is_power_of_two(RAKING_THREADS); + + /// Whether or not accesses into smem are unguarded + static constexpr bool RAKING_UNGUARDED = BlockRakingLayout::UNGUARDED; + + /// Shared memory storage layout type + union _TempStorage + { + /// Storage for warp-synchronous reduction + typename WarpReduce::TempStorage warp_storage; + + /// Padded thread block raking grid + typename BlockRakingLayout::TempStorage raking_grid; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + // Thread fields + _TempStorage& temp_storage; + unsigned int linear_tid; + + /// Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE BlockReduceRaking(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @param[in] reduction_op + * Binary reduction operator + * + * @param[in] partial + * [lane0 only] Warp-wide aggregate reduction of input items + * + * @param[in] num_valid + * Number of valid elements (may be less than BLOCK_THREADS) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T RakingReduction( + ReductionOp reduction_op, T* raking_segment, T partial, int num_valid, constant_t /*iteration*/) + { + // Update partial if addend is in range + if ((IS_FULL_TILE && RAKING_UNGUARDED) || ((linear_tid * SEGMENT_LENGTH) + ITERATION < num_valid)) + { + T addend = raking_segment[ITERATION]; + partial = reduction_op(partial, addend); + } + return RakingReduction(reduction_op, raking_segment, partial, num_valid, constant_t()); + } + + /** + * @param[in] reduction_op + * Binary reduction operator + * + * @param[in] partial + * [lane0 only] Warp-wide aggregate reduction of input items + * + * @param[in] num_valid + * Number of valid elements (may be less than BLOCK_THREADS) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T RakingReduction( + ReductionOp /*reduction_op*/, + T* /*raking_segment*/, + T partial, + int /*num_valid*/, + constant_t /*iteration*/) + { + return partial; + } + + /** + * @brief Computes a thread block-wide reduction using the specified reduction operator. The + * first num_valid threads each contribute one reduction partial. The return value is + * only valid for thread0. + * + * @param[in] partial + * Calling thread's input partial reductions + * + * @param[in] num_valid + * Number of valid elements (may be less than BLOCK_THREADS) + * + * @param[in] reduction_op + * Binary reduction operator + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Reduce(T partial, int num_valid, ReductionOp reduction_op) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp synchronous reduction (unguarded if active threads is a power-of-two) + partial = WarpReduce(temp_storage.warp_storage).template Reduce(partial, num_valid, reduction_op); + } + else + { + // Place partial into shared memory grid. + *BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid) = partial; + + __syncthreads(); + + // Reduce parallelism to one warp + if (linear_tid < RAKING_THREADS) + { + // Raking reduction in grid + T* raking_segment = BlockRakingLayout::RakingPtr(temp_storage.raking_grid, linear_tid); + partial = raking_segment[0]; + + partial = RakingReduction(reduction_op, raking_segment, partial, num_valid, constant_v<1>); + + int valid_raking_threads = (IS_FULL_TILE) ? RAKING_THREADS : (num_valid + SEGMENT_LENGTH - 1) / SEGMENT_LENGTH; + + // sync before re-using shmem (warp_storage/raking_grid are aliased) + static_assert(RAKING_THREADS <= warp_threads, "RAKING_THREADS must be <= warp size."); + unsigned int mask = static_cast((1ull << RAKING_THREADS) - 1); + __syncwarp(mask); + + partial = WarpReduce(temp_storage.warp_storage) + .template Reduce<(IS_FULL_TILE && RAKING_UNGUARDED)>(partial, valid_raking_threads, reduction_op); + } + } + + return partial; + } + + /** + * @brief Computes a thread block-wide reduction using addition (+) as the reduction operator. + * The first num_valid threads each contribute one reduction partial. The return value is + * only valid for thread0. + * + * @param[in] partial + * Calling thread's input partial reductions + * + * @param[in] num_valid + * Number of valid elements (may be less than BLOCK_THREADS) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Sum(T partial, int num_valid) + { + ::cuda::std::plus<> reduction_op; + + return Reduce(partial, num_valid, reduction_op); + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_raking_commutative_only.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_raking_commutative_only.cuh new file mode 100644 index 00000000..ebc37c6e --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_raking_commutative_only.cuh @@ -0,0 +1,207 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::BlockReduceRakingCommutativeOnly provides raking-based methods of parallel reduction across + * a CUDA thread block. Does not support non-commutative reduction operators. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +/** + * @brief BlockReduceRakingCommutativeOnly provides raking-based methods of parallel reduction + * across a CUDA thread block. Does not support non-commutative reduction operators. Does not + * support block sizes that are not a multiple of the warp size. + * + * @tparam T + * Data type being reduced + * + * @tparam BlockDimX + * The thread block length in threads along the X dimension + * + * @tparam BlockDimY + * The thread block length in threads along the Y dimension + * + * @tparam BlockDimZ + * The thread block length in threads along the Z dimension + */ +template +struct BlockReduceRakingCommutativeOnly +{ + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + // The fall-back implementation to use when BLOCK_THREADS is not a multiple of the warp size or not all threads have + // valid values + using FallBack = detail::BlockReduceRaking; + + /// Constants + /// Number of warp threads + static constexpr int WARP_THREADS = warp_threads; + + /// Whether or not to use fall-back + static constexpr bool USE_FALLBACK = ((BLOCK_THREADS % WARP_THREADS != 0) || (BLOCK_THREADS <= WARP_THREADS)); + + /// Number of raking threads + static constexpr int RAKING_THREADS = WARP_THREADS; + + /// Number of threads actually sharing items with the raking threads + static constexpr int SHARING_THREADS = ::cuda::std::max(1, BLOCK_THREADS - RAKING_THREADS); + + /// Number of raking elements per warp synchronous raking thread + static constexpr int SEGMENT_LENGTH = SHARING_THREADS / WARP_THREADS; + + /// WarpReduce utility type + using WarpReduce = WarpReduce; + + /// Layout type for padded thread block raking grid + using BlockRakingLayout = BlockRakingLayout; + + /// Shared memory storage layout type + union _TempStorage + { + struct DefaultStorage + { + /// Storage for warp-synchronous reduction + typename WarpReduce::TempStorage warp_storage; + + /// Padded thread block raking grid + typename BlockRakingLayout::TempStorage raking_grid; + } default_storage; + + /// Fall-back storage for non-commutative block reduction + typename FallBack::TempStorage fallback_storage; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + // Thread fields + _TempStorage& temp_storage; + unsigned int linear_tid; + + /// Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE BlockReduceRakingCommutativeOnly(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + /** + * @brief Computes a thread block-wide reduction using addition (+) as the reduction operator. + * The first num_valid threads each contribute one reduction partial. + * The return value is only valid for thread0. + * + * @param[in] partial + * Calling thread's input partial reductions + * + * @param[in] num_valid + * Number of valid elements (may be less than BLOCK_THREADS) + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Sum(T partial, int num_valid) + { + if (USE_FALLBACK || !FULL_TILE) + { + return FallBack(temp_storage.fallback_storage).template Sum(partial, num_valid); + } + else + { + // Place partial into shared memory grid + if (linear_tid >= RAKING_THREADS) + { + *BlockRakingLayout::PlacementPtr(temp_storage.default_storage.raking_grid, linear_tid - RAKING_THREADS) = + partial; + } + + __syncthreads(); + + // Reduce parallelism to one warp + if (linear_tid < RAKING_THREADS) + { + // Raking reduction in grid + T* raking_segment = BlockRakingLayout::RakingPtr(temp_storage.default_storage.raking_grid, linear_tid); + auto span = ::cuda::std::span(raking_segment, SEGMENT_LENGTH); + partial = cub::ThreadReduce(span, ::cuda::std::plus<>{}, partial); + + // Warp reduction + partial = WarpReduce(temp_storage.default_storage.warp_storage).Sum(partial); + } + } + + return partial; + } + + /** + * @brief Computes a thread block-wide reduction using the specified reduction operator. + * The first num_valid threads each contribute one reduction partial. + * The return value is only valid for thread0. + * + * @param[in] partial + * Calling thread's input partial reductions + * + * @param[in] num_valid + * Number of valid elements (may be less than BLOCK_THREADS) + * + * @param[in] reduction_op + * Binary reduction operator + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Reduce(T partial, int num_valid, ReductionOp reduction_op) + { + if (USE_FALLBACK || !FULL_TILE) + { + return FallBack(temp_storage.fallback_storage).template Reduce(partial, num_valid, reduction_op); + } + else + { + // Place partial into shared memory grid + if (linear_tid >= RAKING_THREADS) + { + *BlockRakingLayout::PlacementPtr(temp_storage.default_storage.raking_grid, linear_tid - RAKING_THREADS) = + partial; + } + + __syncthreads(); + + // Reduce parallelism to one warp + if (linear_tid < RAKING_THREADS) + { + // Raking reduction in grid + T* raking_segment = BlockRakingLayout::RakingPtr(temp_storage.default_storage.raking_grid, linear_tid); + auto span = ::cuda::std::span(raking_segment, SEGMENT_LENGTH); + partial = cub::ThreadReduce(span, reduction_op, partial); + + // Warp reduction + partial = WarpReduce(temp_storage.default_storage.warp_storage).Reduce(partial, reduction_op); + } + } + + return partial; + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_warp_reductions.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_warp_reductions.cuh new file mode 100644 index 00000000..4b23b049 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_reduce_warp_reductions.cuh @@ -0,0 +1,260 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @rst +//! @file +//! cub::BlockReduceWarpReductions provides variants of warp-reduction-based parallel reduction +//! across a CUDA thread block. Supports non-commutative reduction operators. +//! @endrst + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +//! @rst +//! BlockReduceWarpReductions provides variants of warp-reduction-based parallel reduction +//! across a CUDA thread block. Supports non-commutative reduction operators. +//! @endrst +//! +//! @tparam T +//! Data type being reduced +//! +//! @tparam BlockDimX +//! The thread block length in threads along the X dimension +//! +//! @tparam BlockDimY +//! The thread block length in threads along the Y dimension +//! +//! @tparam BlockDimZ +//! The thread block length in threads along the Z dimension +//! +//! @tparam IsDeterministic +//! Whether the reduction is deterministic +template +struct BlockReduceWarpReductions +{ + /// The thread block size in threads + static constexpr int threads_per_block = BlockDimX * BlockDimY * BlockDimZ; + + /// Number of active warps + static constexpr int warps = ::cuda::ceil_div(threads_per_block, warp_threads); + + /// The logical warp size for warp reductions + static constexpr int logical_warp_size = ::cuda::std::min(threads_per_block, warp_threads); + + /// Whether or not the logical warp size evenly divides the thread block size + static constexpr bool even_warp_multiple = (threads_per_block % logical_warp_size == 0); + + using WarpReduceInternal = typename WarpReduce::InternalWarpReduce; + + /// Shared memory storage layout type + struct _TempStorage + { + /// Buffer for warp-synchronous reduction + typename WarpReduceInternal::TempStorage warp_reduce[warps]; + + /// Shared totals from each warp-synchronous reduction + T warp_aggregates[warps]; + + /// Shared prefix for the entire thread block + T block_prefix; + }; + + using TempStorage = Uninitialized<_TempStorage>; + + // Thread fields + _TempStorage& temp_storage; + int linear_tid; + int warp_id; + int lane_id; + + _CCCL_DEVICE _CCCL_FORCEINLINE BlockReduceWarpReductions(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + , warp_id((warps == 1) ? 0 : linear_tid / warp_threads) + , lane_id(static_cast(::cuda::ptx::get_sreg_laneid())) + {} + + //! @rst + //! Returns block-wide aggregate in *thread*\ :sub:`0`. + //! @endrst + //! + //! @tparam ReductionOp + //! **[inferred]** Binary reduction operator type + //! + //! @param[in] reduction_op + //! Binary reduction operator + //! + //! @param[in] warp_aggregate + //! **[**\ *lane*\ :sub:`0` **only]** Warp-wide aggregate reduction of input items + template + _CCCL_DEVICE _CCCL_FORCEINLINE T ApplyWarpAggregatesNonDeterministic(ReductionOp reduction_op, T warp_aggregate) + { + if (linear_tid == 0) + { + detail::uninitialized_copy_single(temp_storage.warp_aggregates, warp_aggregate); + } + + __syncthreads(); + + // Warp 0 already contributed its aggregate above since its also linear_tid == 0 + if (lane_id == 0 && warp_id != 0) + { + // TODO: replace this with other atomic operations when specified + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_60, + ({ + ::cuda::atomic_ref atomic_target(temp_storage.warp_aggregates[0]); + atomic_target.fetch_add(warp_aggregate, ::cuda::memory_order_relaxed); + }), + (atomicAdd(&temp_storage.warp_aggregates[0], warp_aggregate);)); + } + + __syncthreads(); + return temp_storage.warp_aggregates[0]; + } + + //! @rst + //! Recursively applies warp aggregates using template unrolling for deterministic reduction. + //! @endrst + //! + //! @tparam FullTile + //! **[inferred]** Whether this is a full tile + //! + //! @tparam ReductionOp + //! **[inferred]** Binary reduction operator type + template + _CCCL_DEVICE _CCCL_FORCEINLINE T ApplyWarpAggregates(ReductionOp reduction_op, T warp_aggregate, int num_valid) + { + // Share lane aggregates + if (lane_id == 0) + { + detail::uninitialized_copy_single(temp_storage.warp_aggregates + warp_id, warp_aggregate); + } + + __syncthreads(); + + // Update total aggregate in warp 0, lane 0 + if (linear_tid == 0) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int warp_idx = 1; warp_idx < warps; ++warp_idx) + { + if (FullTile || (warp_idx * logical_warp_size < num_valid)) + { + T addend = temp_storage.warp_aggregates[warp_idx]; + warp_aggregate = reduction_op(warp_aggregate, addend); + } + } + } + + return warp_aggregate; + } + + //! @rst + //! Computes a thread block-wide reduction using addition (+/cuda::std::plus<>) as the reduction operator. + //! The first num_valid threads each contribute one reduction partial. The return value is + //! only valid for *thread*\ :sub:`0`. + //! @endrst + //! + //! @tparam FullTile + //! **[inferred]** Whether this is a full tile + //! + //! @param[in] input + //! Calling thread's input partial reductions + //! + //! @param[in] num_valid + //! Number of valid elements (may be less than threads_per_block) + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Sum(T input, int num_valid) + { + ::cuda::std::plus<> reduction_op; + const int warp_offset = (warp_id * logical_warp_size); + const int warp_num_valid = + ((FullTile && even_warp_multiple) || (warp_offset + logical_warp_size <= num_valid)) + ? logical_warp_size + : num_valid - warp_offset; + + // Warp reduction in every warp + T warp_aggregate = WarpReduceInternal(temp_storage.warp_reduce[warp_id]) + .template Reduce<(FullTile && even_warp_multiple)>(input, warp_num_valid, reduction_op); + + // Update outputs and block_aggregate with warp-wide aggregates from lane-0s + if constexpr (IsDeterministic) + { + return ApplyWarpAggregates(reduction_op, warp_aggregate, num_valid); + } + else + { + return ApplyWarpAggregatesNonDeterministic(reduction_op, warp_aggregate); + } + } + + //! @rst + //! Computes a thread block-wide reduction using the specified reduction operator. + //! The first num_valid threads each contribute one reduction partial. + //! The return value is only valid for *thread*\ :sub:`0`. + //! @endrst + //! + //! @tparam FullTile + //! **[inferred]** Whether this is a full tile + //! + //! @tparam ReductionOp + //! **[inferred]** Binary reduction operator type + //! + //! @param[in] input + //! Calling thread's input partial reductions + //! + //! @param[in] num_valid + //! Number of valid elements (may be less than threads_per_block) + //! + //! @param[in] reduction_op + //! Binary reduction operator + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Reduce(T input, int num_valid, ReductionOp reduction_op) + { + const int warp_offset = warp_id * logical_warp_size; + const int warp_num_valid = + ((FullTile && even_warp_multiple) || (warp_offset + logical_warp_size <= num_valid)) + ? logical_warp_size + : num_valid - warp_offset; + + // Warp reduction in every warp + const T warp_aggregate = WarpReduceInternal(temp_storage.warp_reduce[warp_id]) + .template Reduce<(FullTile && even_warp_multiple)>(input, warp_num_valid, reduction_op); + + // Update outputs and block_aggregate with warp-wide aggregates from lane-0s + if constexpr (IsDeterministic) + { + return ApplyWarpAggregates(reduction_op, warp_aggregate, num_valid); + } + else + { + return ApplyWarpAggregatesNonDeterministic(reduction_op, warp_aggregate); + } + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_scan_raking.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_scan_raking.cuh new file mode 100644 index 00000000..c24acd0b --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_scan_raking.cuh @@ -0,0 +1,766 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::BlockScanRaking provides variants of raking-based parallel prefix scan across a + * CUDA thread block. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +/** + * @brief BlockScanRaking provides variants of raking-based parallel prefix scan across a CUDA + * thread block. + * + * @tparam T + * Data type being scanned + * + * @tparam BlockDimX + * The thread block length in threads along the X dimension + * + * @tparam BlockDimY + * The thread block length in threads along the Y dimension + * + * @tparam BlockDimZ + * The thread block length in threads along the Z dimension + * + * @tparam Memoize + * Whether or not to buffer outer raking scan partials to incur fewer shared memory reads at the + * expense of higher register pressure + */ +template +struct BlockScanRaking +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + /// Layout type for padded thread block raking grid + using BlockRakingLayout = BlockRakingLayout; + + /// Constants + /// Number of raking threads + static constexpr int RAKING_THREADS = BlockRakingLayout::RAKING_THREADS; + + /// Number of raking elements per warp synchronous raking thread + static constexpr int SEGMENT_LENGTH = BlockRakingLayout::SEGMENT_LENGTH; + + /// Cooperative work can be entirely warp synchronous + static constexpr bool WARP_SYNCHRONOUS = (BLOCK_THREADS == RAKING_THREADS); + + /// WarpScan utility type + using WarpScan = WarpScan; + + /// Shared memory storage layout type + struct _TempStorage + { + /// Buffer for warp-synchronous scan + typename WarpScan::TempStorage warp_scan; + + /// Padded thread block raking grid + typename BlockRakingLayout::TempStorage raking_grid; + + /// Block aggregate + T block_aggregate; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + // Thread fields + _TempStorage& temp_storage; + unsigned int linear_tid; + T cached_segment[SEGMENT_LENGTH]; + + //--------------------------------------------------------------------- + // Utility methods + //--------------------------------------------------------------------- + + /** + * @brief Templated reduction + * + * @param[in] raking_ptr + * Input array + * + * @param[in] scan_op + * Binary reduction operator + * + * @param[in] raking_partial + * Prefix to seed reduction with + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T + GuardedReduce(T* raking_ptr, ScanOp scan_op, T raking_partial, constant_t /*iteration*/) + { + if ((BlockRakingLayout::UNGUARDED) || (((linear_tid * SEGMENT_LENGTH) + ITERATION) < BLOCK_THREADS)) + { + T addend = raking_ptr[ITERATION]; + raking_partial = scan_op(raking_partial, addend); + } + + return GuardedReduce(raking_ptr, scan_op, raking_partial, constant_v); + } + + /** + * @brief Templated reduction (base case) + * + * @param[in] raking_ptr + * Input array + * + * @param[in] scan_op + * Binary reduction operator + * + * @param[in] raking_partial + * Prefix to seed reduction with + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T + GuardedReduce(T* /*raking_ptr*/, ScanOp /*scan_op*/, T raking_partial, constant_t /*iteration*/) + { + return raking_partial; + } + + /** + * @brief Templated copy + * + * @param out + * [out] Out array + * + * @param in + * [in] Input array + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void CopySegment(T* out, T* in, constant_t /*iteration*/) + { + out[ITERATION] = in[ITERATION]; + CopySegment(out, in, constant_v); + } + + /** + * @brief Templated copy (base case) + * + * @param[out] out + * Out array + * + * @param[in] in + * Input array + */ + _CCCL_DEVICE _CCCL_FORCEINLINE void CopySegment(T* /*out*/, T* /*in*/, constant_t /*iteration*/) {} + + /// Performs upsweep raking reduction, returning the aggregate + template + _CCCL_DEVICE _CCCL_FORCEINLINE T Upsweep(ScanOp scan_op) + { + T* smem_raking_ptr = BlockRakingLayout::RakingPtr(temp_storage.raking_grid, linear_tid); + + // Read data into registers + CopySegment(cached_segment, smem_raking_ptr, constant_v<0>); + + T raking_partial = cached_segment[0]; + + return GuardedReduce(cached_segment, scan_op, raking_partial, constant_v<1>); + } + + /// Performs exclusive downsweep raking scan + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveDownsweep(ScanOp scan_op, T raking_partial, bool apply_prefix = true) + { + T* smem_raking_ptr = BlockRakingLayout::RakingPtr(temp_storage.raking_grid, linear_tid); + + // Read data back into registers + if constexpr (!Memoize) + { + CopySegment(cached_segment, smem_raking_ptr, constant_v<0>); + } + + detail::ThreadScanExclusive(cached_segment, cached_segment, scan_op, raking_partial, apply_prefix); + + // Write data back to smem + CopySegment(smem_raking_ptr, cached_segment, constant_v<0>); + } + + /// Performs inclusive downsweep raking scan + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveDownsweep(ScanOp scan_op, T raking_partial, bool apply_prefix = true) + { + T* smem_raking_ptr = BlockRakingLayout::RakingPtr(temp_storage.raking_grid, linear_tid); + + // Read data back into registers + if constexpr (!Memoize) + { + CopySegment(cached_segment, smem_raking_ptr, constant_v<0>); + } + + detail::ThreadScanInclusive(cached_segment, cached_segment, scan_op, raking_partial, apply_prefix); + + // Write data back to smem + CopySegment(smem_raking_ptr, cached_segment, constant_v<0>); + } + + //--------------------------------------------------------------------- + // Constructors + //--------------------------------------------------------------------- + + /// Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE BlockScanRaking(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + {} + + //--------------------------------------------------------------------- + // Exclusive scans + //--------------------------------------------------------------------- + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. With no initial value, + * the output computed for thread0 is undefined. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] exclusive_output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& exclusive_output, ScanOp scan_op) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp-synchronous scan + WarpScan(temp_storage.warp_scan).ExclusiveScan(input, exclusive_output, scan_op); + } + else + { + // Place thread partial into shared memory raking grid + T* placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid); + detail::uninitialized_copy_single(placement_ptr, input); + + __syncthreads(); + + // Reduce parallelism down to just raking threads + if (linear_tid < RAKING_THREADS) + { + // Raking upsweep reduction across shared partials + T upsweep_partial = Upsweep(scan_op); + + // Warp-synchronous scan + T exclusive_partial; + WarpScan(temp_storage.warp_scan).ExclusiveScan(upsweep_partial, exclusive_partial, scan_op); + + // Exclusive raking downsweep scan + ExclusiveDownsweep(scan_op, exclusive_partial, (linear_tid != 0)); + } + + __syncthreads(); + + // Grab thread prefix from shared memory + exclusive_output = *placement_ptr; + } + } + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. + * + * @param[in] input + * Calling thread's input items + * + * @param[out] output + * Calling thread's output items (may be aliased to \p input) + * + * @param[in] initial_value + * Initial value to seed the exclusive scan + * + * @param[in] scan_op + * Binary scan operator + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, const T& initial_value, ScanOp scan_op) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp-synchronous scan + WarpScan(temp_storage.warp_scan).ExclusiveScan(input, output, initial_value, scan_op); + } + else + { + // Place thread partial into shared memory raking grid + T* placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid); + detail::uninitialized_copy_single(placement_ptr, input); + + __syncthreads(); + + // Reduce parallelism down to just raking threads + if (linear_tid < RAKING_THREADS) + { + // Raking upsweep reduction across shared partials + T upsweep_partial = Upsweep(scan_op); + + // Exclusive Warp-synchronous scan + T exclusive_partial; + WarpScan(temp_storage.warp_scan).ExclusiveScan(upsweep_partial, exclusive_partial, initial_value, scan_op); + + // Exclusive raking downsweep scan + ExclusiveDownsweep(scan_op, exclusive_partial); + } + + __syncthreads(); + + // Grab exclusive partial from shared memory + output = *placement_ptr; + } + } + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. Also provides every + * thread with the block-wide \p block_aggregate of all inputs. With no initial value, + * the output computed for thread0 is undefined. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, ScanOp scan_op, T& block_aggregate) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp-synchronous scan + WarpScan(temp_storage.warp_scan).ExclusiveScan(input, output, scan_op, block_aggregate); + } + else + { + // Place thread partial into shared memory raking grid + T* placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid); + detail::uninitialized_copy_single(placement_ptr, input); + + __syncthreads(); + + // Reduce parallelism down to just raking threads + if (linear_tid < RAKING_THREADS) + { + // Raking upsweep reduction across shared partials + T upsweep_partial = Upsweep(scan_op); + + // Warp-synchronous scan + T inclusive_partial; + T exclusive_partial; + WarpScan(temp_storage.warp_scan).Scan(upsweep_partial, inclusive_partial, exclusive_partial, scan_op); + + // Exclusive raking downsweep scan + ExclusiveDownsweep(scan_op, exclusive_partial, (linear_tid != 0)); + + // Broadcast aggregate to all threads + if (linear_tid == RAKING_THREADS - 1) + { + temp_storage.block_aggregate = inclusive_partial; + } + } + + __syncthreads(); + + // Grab thread prefix from shared memory + output = *placement_ptr; + + // Retrieve block aggregate + block_aggregate = temp_storage.block_aggregate; + } + } + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. Also provides every + * thread with the block-wide \p block_aggregate of all inputs. + * + * @param[in] input + * Calling thread's input items + * + * @param[out] output + * Calling thread's output items (may be aliased to \p input) + * + * @param[in] initial_value + * Initial value to seed the exclusive scan + * + * @param[in] scan_op + * Binary scan operator + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T input, T& output, const T& initial_value, ScanOp scan_op, T& block_aggregate) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp-synchronous scan + WarpScan(temp_storage.warp_scan).ExclusiveScan(input, output, initial_value, scan_op, block_aggregate); + } + else + { + // Place thread partial into shared memory raking grid + T* placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid); + detail::uninitialized_copy_single(placement_ptr, input); + + __syncthreads(); + + // Reduce parallelism down to just raking threads + if (linear_tid < RAKING_THREADS) + { + // Raking upsweep reduction across shared partials + T upsweep_partial = Upsweep(scan_op); + + // Warp-synchronous scan + T exclusive_partial; + WarpScan(temp_storage.warp_scan) + .ExclusiveScan(upsweep_partial, exclusive_partial, initial_value, scan_op, block_aggregate); + + // Exclusive raking downsweep scan + ExclusiveDownsweep(scan_op, exclusive_partial); + + // Broadcast aggregate to other threads + if (linear_tid == 0) + { + temp_storage.block_aggregate = block_aggregate; + } + } + + __syncthreads(); + + // Grab exclusive partial from shared memory + output = *placement_ptr; + + // Retrieve block aggregate + block_aggregate = temp_storage.block_aggregate; + } + } + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. the call-back functor \p + * block_prefix_callback_op is invoked by the first warp in the block, and the value + * returned by lane0 in that warp is used as the "seed" value that + * logically prefixes the thread block's scan inputs. Also provides every thread with + * the block-wide \p block_aggregate of all inputs. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + * + * @param[in-out] block_prefix_callback_op + * [warp0 only] Call-back functor for specifying a thread + * block-wide prefix to be applied to all inputs. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T input, T& output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp-synchronous scan + T block_aggregate; + WarpScan warp_scan(temp_storage.warp_scan); + warp_scan.ExclusiveScan(input, output, scan_op, block_aggregate); + + // Obtain warp-wide prefix in lane0, then broadcast to other lanes + T block_prefix = block_prefix_callback_op(block_aggregate); + block_prefix = warp_scan.Broadcast(block_prefix, 0); + + output = scan_op(block_prefix, output); + if (linear_tid == 0) + { + output = block_prefix; + } + } + else + { + // Place thread partial into shared memory raking grid + T* placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid); + detail::uninitialized_copy_single(placement_ptr, input); + + __syncthreads(); + + // Reduce parallelism down to just raking threads + if (linear_tid < RAKING_THREADS) + { + WarpScan warp_scan(temp_storage.warp_scan); + + // Raking upsweep reduction across shared partials + T upsweep_partial = Upsweep(scan_op); + + // Warp-synchronous scan + T exclusive_partial, block_aggregate; + warp_scan.ExclusiveScan(upsweep_partial, exclusive_partial, scan_op, block_aggregate); + + // Obtain block-wide prefix in lane0, then broadcast to other lanes + T block_prefix = block_prefix_callback_op(block_aggregate); + block_prefix = warp_scan.Broadcast(block_prefix, 0); + + // Update prefix with warpscan exclusive partial + T downsweep_prefix = scan_op(block_prefix, exclusive_partial); + if (linear_tid == 0) + { + downsweep_prefix = block_prefix; + } + + // Exclusive raking downsweep scan + ExclusiveDownsweep(scan_op, downsweep_prefix); + } + + __syncthreads(); + + // Grab thread prefix from shared memory + output = *placement_ptr; + } + } + + //--------------------------------------------------------------------- + // Inclusive scans + //--------------------------------------------------------------------- + + /** + * @brief Computes an inclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& output, ScanOp scan_op) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp-synchronous scan + WarpScan(temp_storage.warp_scan).InclusiveScan(input, output, scan_op); + } + else + { + // Place thread partial into shared memory raking grid + T* placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid); + detail::uninitialized_copy_single(placement_ptr, input); + + __syncthreads(); + + // Reduce parallelism down to just raking threads + if (linear_tid < RAKING_THREADS) + { + // Raking upsweep reduction across shared partials + T upsweep_partial = Upsweep(scan_op); + + // Exclusive Warp-synchronous scan + T exclusive_partial; + WarpScan(temp_storage.warp_scan).ExclusiveScan(upsweep_partial, exclusive_partial, scan_op); + + // Inclusive raking downsweep scan + InclusiveDownsweep(scan_op, exclusive_partial, (linear_tid != 0)); + } + + __syncthreads(); + + // Grab thread prefix from shared memory + output = *placement_ptr; + } + } + + /** + * @brief Computes an inclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. Also provides every + * thread with the block-wide \p block_aggregate of all inputs. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& output, ScanOp scan_op, T& block_aggregate) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp-synchronous scan + WarpScan(temp_storage.warp_scan).InclusiveScan(input, output, scan_op, block_aggregate); + } + else + { + // Place thread partial into shared memory raking grid + T* placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid); + detail::uninitialized_copy_single(placement_ptr, input); + + __syncthreads(); + + // Reduce parallelism down to just raking threads + if (linear_tid < RAKING_THREADS) + { + // Raking upsweep reduction across shared partials + T upsweep_partial = Upsweep(scan_op); + + // Warp-synchronous scan + T inclusive_partial; + T exclusive_partial; + WarpScan(temp_storage.warp_scan).Scan(upsweep_partial, inclusive_partial, exclusive_partial, scan_op); + + // Inclusive raking downsweep scan + InclusiveDownsweep(scan_op, exclusive_partial, (linear_tid != 0)); + + // Broadcast aggregate to all threads + if (linear_tid == RAKING_THREADS - 1) + { + temp_storage.block_aggregate = inclusive_partial; + } + } + + __syncthreads(); + + // Grab thread prefix from shared memory + output = *placement_ptr; + + // Retrieve block aggregate + block_aggregate = temp_storage.block_aggregate; + } + } + + /** + * @brief Computes an inclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. the call-back functor \p + * block_prefix_callback_op is invoked by the first warp in the block, and the value + * returned by lane0 in that warp is used as the "seed" value that + * logically prefixes the thread block's scan inputs. Also provides every thread with + * the block-wide \p block_aggregate of all inputs. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + * + * @param[in-out] block_prefix_callback_op + * [warp0 only] Call-back functor for specifying a thread + * block-wide prefix to be applied to all inputs. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + InclusiveScan(T input, T& output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op) + { + if (WARP_SYNCHRONOUS) + { + // Short-circuit directly to warp-synchronous scan + T block_aggregate; + WarpScan warp_scan(temp_storage.warp_scan); + warp_scan.InclusiveScan(input, output, scan_op, block_aggregate); + + // Obtain warp-wide prefix in lane0, then broadcast to other lanes + T block_prefix = block_prefix_callback_op(block_aggregate); + block_prefix = warp_scan.Broadcast(block_prefix, 0); + + // Update prefix with exclusive warpscan partial + output = scan_op(block_prefix, output); + } + else + { + // Place thread partial into shared memory raking grid + T* placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid); + detail::uninitialized_copy_single(placement_ptr, input); + + __syncthreads(); + + // Reduce parallelism down to just raking threads + if (linear_tid < RAKING_THREADS) + { + WarpScan warp_scan(temp_storage.warp_scan); + + // Raking upsweep reduction across shared partials + T upsweep_partial = Upsweep(scan_op); + + // Warp-synchronous scan + T exclusive_partial, block_aggregate; + warp_scan.ExclusiveScan(upsweep_partial, exclusive_partial, scan_op, block_aggregate); + + // Obtain block-wide prefix in lane0, then broadcast to other lanes + T block_prefix = block_prefix_callback_op(block_aggregate); + block_prefix = warp_scan.Broadcast(block_prefix, 0); + + // Update prefix with warpscan exclusive partial + T downsweep_prefix = scan_op(block_prefix, exclusive_partial); + if (linear_tid == 0) + { + downsweep_prefix = block_prefix; + } + + // Inclusive raking downsweep scan + InclusiveDownsweep(scan_op, downsweep_prefix); + } + + __syncthreads(); + + // Grab thread prefix from shared memory + output = *placement_ptr; + } + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_scan_warp_scans.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_scan_warp_scans.cuh new file mode 100644 index 00000000..12f7127d --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_scan_warp_scans.cuh @@ -0,0 +1,514 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::BlockScanWarpscans provides warpscan-based variants of parallel prefix scan across a CUDA thread block. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +/** + * @brief BlockScanWarpScans provides warpscan-based variants of parallel prefix scan across a CUDA + * thread block. + * + * @tparam BlockDimX + * The thread block length in threads along the X dimension + * + * @tparam BlockDimY + * The thread block length in threads along the Y dimension + * + * @tparam BlockDimZ + * The thread block length in threads along the Z dimension + */ +template +struct BlockScanWarpScans +{ + //--------------------------------------------------------------------- + // Types and constants + //--------------------------------------------------------------------- + + /// Constants + /// Number of warp threads + static constexpr int WARP_THREADS = warp_threads; + + /// The thread block size in threads + static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ; + + /// Number of active warps + static constexpr int WARPS = ::cuda::ceil_div(BLOCK_THREADS, WARP_THREADS); + + /// WarpScan utility type + using WarpScanT = WarpScan; + + /// WarpScan utility type + using WarpAggregateScan = WarpScan; + + /// Shared memory storage layout type + + struct __align__(32) _TempStorage + { + T warp_aggregates[WARPS]; + + /// Buffer for warp-synchronous scans + typename WarpScanT::TempStorage warp_scan[WARPS]; + + /// Shared prefix for the entire thread block + T block_prefix; + }; + + /// Alias wrapper allowing storage to be unioned + struct TempStorage : Uninitialized<_TempStorage> + {}; + + //--------------------------------------------------------------------- + // Per-thread fields + //--------------------------------------------------------------------- + + // Thread fields + _TempStorage& temp_storage; + unsigned int linear_tid; + unsigned int warp_id; + unsigned int lane_id; + + //--------------------------------------------------------------------- + // Constructors + //--------------------------------------------------------------------- + + /// Constructor + _CCCL_DEVICE _CCCL_FORCEINLINE BlockScanWarpScans(TempStorage& temp_storage) + : temp_storage(temp_storage.Alias()) + , linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ)) + , warp_id((WARPS == 1) ? 0 : linear_tid / WARP_THREADS) + , lane_id(::cuda::ptx::get_sreg_laneid()) + {} + + //--------------------------------------------------------------------- + // Utility methods + //--------------------------------------------------------------------- + + /** + * @param[out] warp_prefix + * The calling thread's partial reduction + * + * @param[in] scan_op + * Binary scan operator + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ApplyWarpAggregates(T& warp_prefix, ScanOp scan_op, T& block_aggregate, constant_t /*addend_warp*/) + { + if (warp_id == WARP) + { + warp_prefix = block_aggregate; + } + + T addend = temp_storage.warp_aggregates[WARP]; + block_aggregate = scan_op(block_aggregate, addend); + + ApplyWarpAggregates(warp_prefix, scan_op, block_aggregate, constant_v); + } + + /** + * @param[out] warp_prefix + * The calling thread's partial reduction + * + * @param[in] scan_op + * Binary scan operator + * + * @param[out] block_aggregat + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ApplyWarpAggregates(T& /*warp_prefix*/, ScanOp /*scan_op*/, T& /*block_aggregate*/, constant_t /*addend_warp*/) + {} + + /** + * @brief Use the warp-wide aggregates to compute the calling warp's prefix. Also returns + * block-wide aggregate in all threads. + * + * @param[in] scan_op + * Binary scan operator + * + * @param[in] warp_aggregate + * [laneWARP_THREADS - 1 only] Warp-wide aggregate reduction of + * input items + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T ComputeWarpPrefix(ScanOp scan_op, T warp_aggregate, T& block_aggregate) + { + // Last lane in each warp shares its warp-aggregate + if (lane_id == WARP_THREADS - 1) + { + detail::uninitialized_copy_single(temp_storage.warp_aggregates + warp_id, warp_aggregate); + } + + __syncthreads(); + + // Accumulate block aggregates and save the one that is our warp's prefix + T warp_prefix; + block_aggregate = temp_storage.warp_aggregates[0]; + + // Use template unrolling (since the PTX backend can't handle unrolling it for SM1x) + // TODO(bgruber): does that still hold today? This is creating a lot of template instantiations + ApplyWarpAggregates(warp_prefix, scan_op, block_aggregate, constant_v<1>); + /* + _CCCL_PRAGMA_UNROLL_FULL() + for (int WARP = 1; WARP < WARPS; ++WARP) + { + if (warp_id == WARP) + warp_prefix = block_aggregate; + + T addend = temp_storage.warp_aggregates[WARP]; + block_aggregate = scan_op(block_aggregate, addend); + } + */ + + return warp_prefix; + } + + /** + * @brief Use the warp-wide aggregates and initial-value to compute the calling warp's prefix. + * Also returns block-wide aggregate in all threads. + * + * @param[in] scan_op + * Binary scan operator + * + * @param[in] warp_aggregate + * [laneWARP_THREADS - 1 only] Warp-wide aggregate reduction of + * input items + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + * + * @param[in] initial_value + * Initial value to seed the exclusive scan + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE T + ComputeWarpPrefix(ScanOp scan_op, T warp_aggregate, T& block_aggregate, const T& initial_value) + { + T warp_prefix = ComputeWarpPrefix(scan_op, warp_aggregate, block_aggregate); + + warp_prefix = scan_op(initial_value, warp_prefix); + + if (warp_id == 0) + { + warp_prefix = initial_value; + } + + return warp_prefix; + } + + //--------------------------------------------------------------------- + // Exclusive scans + //--------------------------------------------------------------------- + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. With no initial value, + * the output computed for thread0 is undefined. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] exclusive_output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& exclusive_output, ScanOp scan_op) + { + // Compute block-wide exclusive scan. The exclusive output from tid0 is invalid. + T block_aggregate; + ExclusiveScan(input, exclusive_output, scan_op, block_aggregate); + } + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. + * + * @param[in] input + * Calling thread's input items + * + * @param[out] exclusive_output + * Calling thread's output items (may be aliased to \p input) + * + * @param[in] initial_value + * Initial value to seed the exclusive scan + * + * @param[in] scan_op + * Binary scan operator + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& exclusive_output, const T& initial_value, ScanOp scan_op) + { + T block_aggregate; + ExclusiveScan(input, exclusive_output, initial_value, scan_op, block_aggregate); + } + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. Also provides every + * thread with the block-wide \p block_aggregate of all inputs. With no initial value, + * the output computed for thread0 is undefined. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] exclusive_output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& exclusive_output, ScanOp scan_op, T& block_aggregate) + { + // Compute warp scan in each warp. The exclusive output from each lane0 is invalid. + T inclusive_output; + WarpScanT(temp_storage.warp_scan[warp_id]).Scan(input, inclusive_output, exclusive_output, scan_op); + + // Compute the warp-wide prefix and block-wide aggregate for each warp. Warp prefix for warp0 is invalid. + T warp_prefix = ComputeWarpPrefix(scan_op, inclusive_output, block_aggregate); + + // Apply warp prefix to our lane's partial + if (warp_id != 0) + { + exclusive_output = scan_op(warp_prefix, exclusive_output); + if (lane_id == 0) + { + exclusive_output = warp_prefix; + } + } + } + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. Also provides every + * thread with the block-wide \p block_aggregate of all inputs. + * + * @param[in] input + * Calling thread's input items + * + * @param[out] exclusive_output + * Calling thread's output items (may be aliased to \p input) + * + * @param[in] initial_value + * Initial value to seed the exclusive scan + * + * @param[in] scan_op + * Binary scan operator + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T input, T& exclusive_output, const T& initial_value, ScanOp scan_op, T& block_aggregate) + { + // Compute warp scan in each warp. The exclusive output from each lane0 is invalid. + T inclusive_output; + WarpScanT(temp_storage.warp_scan[warp_id]).Scan(input, inclusive_output, exclusive_output, scan_op); + + // Compute the warp-wide prefix and block-wide aggregate for each warp + T warp_prefix = ComputeWarpPrefix(scan_op, inclusive_output, block_aggregate, initial_value); + + // Apply warp prefix to our lane's partial + exclusive_output = scan_op(warp_prefix, exclusive_output); + if (lane_id == 0) + { + exclusive_output = warp_prefix; + } + } + + /** + * @brief Computes an exclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. the call-back functor \p + * block_prefix_callback_op is invoked by the first warp in the block, and the value + * returned by lane0 in that warp is used as the "seed" value that + * logically prefixes the thread block's scan inputs. Also provides every thread with + * the block-wide \p block_aggregate of all inputs. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] exclusive_output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + * + * @param[in-out] block_prefix_callback_op + * [warp0 only] Call-back functor for specifying a thread + * block-wide prefix to be applied to all inputs. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + ExclusiveScan(T input, T& exclusive_output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op) + { + // Compute block-wide exclusive scan. The exclusive output from tid0 is invalid. + T block_aggregate; + ExclusiveScan(input, exclusive_output, scan_op, block_aggregate); + + // Use the first warp to determine the thread block prefix, returning the result in lane0 + if (warp_id == 0) + { + T block_prefix = block_prefix_callback_op(block_aggregate); + if (lane_id == 0) + { + // Share the prefix with all threads + detail::uninitialized_copy_single(&temp_storage.block_prefix, block_prefix); + + exclusive_output = block_prefix; // The block prefix is the exclusive output for tid0 + } + } + + __syncthreads(); + + // Incorporate thread block prefix into outputs + T block_prefix = temp_storage.block_prefix; + if (linear_tid > 0) + { + exclusive_output = scan_op(block_prefix, exclusive_output); + } + } + + //--------------------------------------------------------------------- + // Inclusive scans + //--------------------------------------------------------------------- + + /** + * @brief Computes an inclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] inclusive_output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& inclusive_output, ScanOp scan_op) + { + T block_aggregate; + InclusiveScan(input, inclusive_output, scan_op, block_aggregate); + } + + /** + * @brief Computes an inclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. Also provides every + * thread with the block-wide \p block_aggregate of all inputs. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] inclusive_output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + * + * @param[out] block_aggregate + * Threadblock-wide aggregate reduction of input items + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& inclusive_output, ScanOp scan_op, T& block_aggregate) + { + WarpScanT(temp_storage.warp_scan[warp_id]).InclusiveScan(input, inclusive_output, scan_op); + + // Compute the warp-wide prefix and block-wide aggregate for each warp. Warp prefix for warp0 is invalid. + T warp_prefix = ComputeWarpPrefix(scan_op, inclusive_output, block_aggregate); + + // Apply warp prefix to our lane's partial + if (warp_id != 0) + { + inclusive_output = scan_op(warp_prefix, inclusive_output); + } + } + + /** + * @brief Computes an inclusive thread block-wide prefix scan using the specified binary \p + * scan_op functor. Each thread contributes one input element. the call-back functor \p + * block_prefix_callback_op is invoked by the first warp in the block, and the value + * returned by lane0 in that warp is used as the "seed" value that + * logically prefixes the thread block's scan inputs. Also provides every thread with + * the block-wide \p block_aggregate of all inputs. + * + * @param[in] input + * Calling thread's input item + * + * @param[out] exclusive_output + * Calling thread's output item (may be aliased to \p input) + * + * @param[in] scan_op + * Binary scan operator + * + * @param[in-out] block_prefix_callback_op + * [warp0 only] Call-back functor for specifying a thread + * block-wide prefix to be applied to all inputs. + */ + template + _CCCL_DEVICE _CCCL_FORCEINLINE void + InclusiveScan(T input, T& exclusive_output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op) + { + T block_aggregate; + InclusiveScan(input, exclusive_output, scan_op, block_aggregate); + + // Use the first warp to determine the thread block prefix, returning the result in lane0 + if (warp_id == 0) + { + T block_prefix = block_prefix_callback_op(block_aggregate); + if (lane_id == 0) + { + // Share the prefix with all threads + detail::uninitialized_copy_single(&temp_storage.block_prefix, block_prefix); + } + } + + __syncthreads(); + + // Incorporate thread block prefix into outputs + T block_prefix = temp_storage.block_prefix; + exclusive_output = scan_op(block_prefix, exclusive_output); + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_topk_air.cuh b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_topk_air.cuh new file mode 100644 index 00000000..b565ec5d --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/block/specializations/block_topk_air.cuh @@ -0,0 +1,523 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +template +struct compare_key_prefix_op +{ + static_assert(::cuda::std::is_unsigned_v, "SortKeyT must be an unsigned type"); + + SortKeyT prefix_mask; + SortKeyT key_prefix; + [[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE constexpr bool operator()(SortKeyT sort_key) const noexcept + { + return (sort_key & prefix_mask) == (key_prefix); + } +}; + +//! @brief Block-level top-k by radix selection. +//! +//! Selects the smallest (or largest) @p k keys from a tile of keys in registers, without +//! fully sorting. The algorithm has two stages: (1) Radix selection determines the bit-prefix +//! of the k-th key by processing bits MSB to LSB in passes of @p RadixBits. In each pass, a +//! histogram over the current digit is built over candidates only (keys matching the prefix so +//! far), then a prefix sum identifies the bucket containing the k-th item. Items in earlier +//! buckets are guaranteed top-k; items in later buckets are discarded; the chosen bucket +//! becomes the candidate set for the next pass. No data movement occurs during this stage—only +//! the histogram in shared memory is updated. (2) Partitioning scatters the top-k items (key +//! prefix <= k-th prefix) into shared memory via atomic counters, then each thread reads back +//! its portion. Supports key-only and key-value selection. +template +class block_topk_air +{ +private: + // TODO (elstehle): Make this configurable + // Whether to include all items tied with the k-th key when selecting top-k + static constexpr bool expand_k_to_include_ties = false; + + static constexpr int threads_per_block = ThreadsPerBlock; + static constexpr int items_per_thread = ItemsPerThread; + static constexpr int tile_items = threads_per_block * items_per_thread; + static constexpr int num_buckets = int{1u << RadixBits}; + + // Calculate number of buckets processed per thread + static constexpr int buckets_per_thread = ::cuda::ceil_div(num_buckets, threads_per_block); + static constexpr bool keys_only = ::cuda::std::is_same_v; + + using histo_counter_t = ::cuda::std::uint32_t; + using block_scan_t = BlockScan; + + using traits = detail::radix::traits_t; + using bit_ordered_type = typename traits::bit_ordered_type; + using bit_ordered_conversion = typename traits::bit_ordered_conversion_policy; + + using fundamental_digit_extractor_t = BFEDigitExtractor; + + struct TempStorage_ + { + union + { + struct + { + histo_counter_t histogram[num_buckets]; + typename block_scan_t::TempStorage scan_temp_storage; + struct + { + histo_counter_t selected; + histo_counter_t candidates; + int bucket; + } pass_state; + } passes; + + struct + { + histo_counter_t selected_offset[2]; + union + { + KeyT keys[tile_items]; + ValueT values[tile_items]; + } exchange; + } select; + } stage; + }; + + /// Shared storage reference + TempStorage_& storage; + + /// Linear thread index + int linear_tid; + + // Initialize histogram bins to zero + _CCCL_DEVICE_API _CCCL_FORCEINLINE void init_histograms() + { + // Initialize histogram bin counts to zeros + int histo_offset = 0; + + // Loop unrolling is beneficial for performance here + _CCCL_PRAGMA_UNROLL_FULL() + for (; histo_offset + threads_per_block <= num_buckets; histo_offset += threads_per_block) + { + storage.stage.passes.histogram[histo_offset + threadIdx.x] = 0; + } + // Finish up with guarded initialization if necessary + if ((num_buckets % threads_per_block != 0) && (histo_offset + threadIdx.x < num_buckets)) + { + storage.stage.passes.histogram[histo_offset + threadIdx.x] = 0; + } + } + + // Compute histogram over keys + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void compute_histograms( + const bit_ordered_type (&unsigned_keys)[items_per_thread], + int valid_items, + DigitExtractorT digit_extractor, + FilterOpT filter_op) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items_per_thread; ++i) + { + const auto item_index = linear_tid * items_per_thread + i; + const bit_ordered_type key = unsigned_keys[i]; + if ((IsFullTile || item_index < valid_items) && filter_op(key)) + { + const auto digit = static_cast(digit_extractor.Digit(key)); + const auto bucket = (SelectDirection == detail::topk::select::min) ? digit : (num_buckets - 1 - digit); + atomicAdd(&storage.stage.passes.histogram[bucket], histo_counter_t{1}); + } + } + } + + // Compute prefix sum over buckets + _CCCL_DEVICE_API _CCCL_FORCEINLINE void compute_bin_offsets() + { + histo_counter_t thread_buckets[buckets_per_thread]{}; + const int base = linear_tid * buckets_per_thread; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < buckets_per_thread; ++i) + { + const int bin_idx = base + i; + if (bin_idx < num_buckets) + { + thread_buckets[i] = storage.stage.passes.histogram[bin_idx]; + } + } + + block_scan_t(storage.stage.passes.scan_temp_storage).InclusiveSum(thread_buckets, thread_buckets); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < buckets_per_thread; ++i) + { + const int bin_idx = base + i; + if (bin_idx < num_buckets) + { + storage.stage.passes.histogram[bin_idx] = thread_buckets[i]; + } + } + } + + // Identify the bucket that the k-th item falls into + _CCCL_DEVICE_API _CCCL_FORCEINLINE void choose_bucket(histo_counter_t k) + { + const int base = linear_tid * buckets_per_thread; + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < buckets_per_thread; ++i) + { + const int bin_idx = base + i; + if (bin_idx < num_buckets) + { + const histo_counter_t prev = (bin_idx == 0) ? 0 : storage.stage.passes.histogram[bin_idx - 1]; + const histo_counter_t cur = storage.stage.passes.histogram[bin_idx]; + + if (prev < k && cur >= k) + { + storage.stage.passes.pass_state.bucket = bin_idx; + storage.stage.passes.pass_state.candidates = cur - prev; + storage.stage.passes.pass_state.selected = prev; + } + } + } + } + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void get_kth_key_prefix( + bit_ordered_type (&unsigned_keys)[items_per_thread], + int k, + int valid_items, + int begin_bit, + int end_bit, + int& total_selected, + int& num_candidates, + bit_ordered_type& kth_key_prefix, + bit_ordered_type& prefix_mask, + DecomposerT decomposer = DecomposerT{}) + { + // Preconditions + [[maybe_unused]] constexpr int max_bit = int(sizeof(KeyT) * 8); + _CCCL_ASSERT(k > 0 && k <= tile_items, "k must be in (0, tile_items]"); + if constexpr (!IsFullTile) + { + _CCCL_ASSERT(valid_items > 0 && valid_items <= tile_items, "valid_items must be in [1, tile_items]"); + } + _CCCL_ASSERT(begin_bit >= 0 && begin_bit < max_bit, "begin_bit must be in [0, max_bit)"); + _CCCL_ASSERT(end_bit > begin_bit && end_bit <= max_bit, "end_bit must be in (begin_bit, max_bit]"); + + // We only consider candidates identified in the previous pass, i.e., ((sortkey & prefix_mask) == kth_prefix) + // With each pass, we identify a wider prefix of the splitter key + kth_key_prefix = 0; + prefix_mask = 0; + + // The total number of selected items + total_selected = 0; + + const int total_bits = (::cuda::std::max) (end_bit - begin_bit, 0); + const int num_passes = ::cuda::ceil_div(total_bits, RadixBits); + for (int pass = 0; pass < num_passes; ++pass) + { + // Bit-range & mask of the current pass + const int pass_end_bit = end_bit - pass * RadixBits; + const int pass_begin_bit = (::cuda::std::max) (pass_end_bit - RadixBits, begin_bit); + const int pass_bits = pass_end_bit - pass_begin_bit; + const bit_ordered_type pass_mask = ::cuda::bitmask(pass_begin_bit, pass_bits); + + // Zero-initialize histograms for the current pass + init_histograms(); + __syncthreads(); + + // Compute histogram over the current pass's, bits pre-filtered for keys matching the previous pass's prefix mask + auto filter_op = compare_key_prefix_op{prefix_mask, kth_key_prefix}; + auto digit_extractor = + traits::template digit_extractor(pass_begin_bit, pass_bits, decomposer); + compute_histograms(unsigned_keys, valid_items, digit_extractor, filter_op); + __syncthreads(); + + // Compute prefix sum over buckets + compute_bin_offsets(); + __syncthreads(); + + // Identify the bucket that the k-th item falls into + choose_bucket(k); + __syncthreads(); + + // Update the current k and length for the next pass + k -= storage.stage.passes.pass_state.selected; + num_candidates = storage.stage.passes.pass_state.candidates; + total_selected += storage.stage.passes.pass_state.selected; + + // Update the kth_key_prefix and prefix_mask for the next pass + // Basically, we will have valid_items candidates with the prefix kth_key_prefix + const auto kth_key_digit = + (SelectDirection == detail::topk::select::min) + ? storage.stage.passes.pass_state.bucket + : (num_buckets - 1 - storage.stage.passes.pass_state.bucket); + kth_key_prefix |= bit_ordered_type(kth_key_digit) << pass_begin_bit; + prefix_mask |= pass_mask; + + // Short-circuit if all candidates are amongst the top-k + if (num_candidates == k) + { + break; + } + } + + // Ensure we can repurpose shared memory after the multi-pass stage + __syncthreads(); + } + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void select_topk( + KeyT (&keys)[items_per_thread], + ValueT (&values)[items_per_thread], + int k, + int valid_items, + int begin_bit, + int end_bit) + { + if constexpr (!IsFullTile) + { + _CCCL_ASSERT(valid_items > 0 && valid_items <= tile_items, "valid_items must be in [1, tile_items]"); + } + + // TODO (elstehle): Short-circuit if k is constrained to be positive + if (k <= 0) + { + return; + } + + // TODO (elstehle): Short-circuit if begin_bit is constrained to be non-negative + begin_bit = (::cuda::std::max) (begin_bit, 0); + + // TODO (elstehle): Short-circuit if end_bit is constrained to be less than the maximum number of bits in the key + // type + const int max_bit = int(sizeof(KeyT) * 8); + if (end_bit > max_bit) + { + end_bit = max_bit; + } + + // TODO (elstehle): Short-circuit if k is greater than the number of items in the tile + if ((!IsFullTile && k >= valid_items) || k >= tile_items) + { + return; + } + + // TODO (elstehle): Add support for custom decomposers + identity_decomposer_t decomposer; + + // Get bit-twiddled sortkeys. For float keys, track which were -0.0 (normalized to +0.0 for ranking) so we can + // restore -0.0 in the output via a bitvector; no extra key buffer. + bit_ordered_type(&unsigned_keys)[ItemsPerThread] = reinterpret_cast(keys); + constexpr int flip_back_num_words = ::cuda::ceil_div(items_per_thread, 32); + [[maybe_unused]] ::cuda::std::uint32_t flip_back_bits[flip_back_num_words] = {}; + if constexpr (::cuda::is_floating_point_v) + { + const bit_ordered_type twiddled_minus_zero = + Traits::TwiddleIn(bit_ordered_type(1) << (8 * sizeof(bit_ordered_type) - 1)); + const bit_ordered_type twiddled_zero = Traits::TwiddleIn(0); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items_per_thread; ++i) + { + unsigned_keys[i] = bit_ordered_conversion::to_bit_ordered(decomposer, unsigned_keys[i]); + if (unsigned_keys[i] == twiddled_minus_zero) + { + flip_back_bits[i / 32] |= (1u << (i % 32)); + unsigned_keys[i] = twiddled_zero; + } + } + } + else + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items_per_thread; ++i) + { + unsigned_keys[i] = bit_ordered_conversion::to_bit_ordered(decomposer, unsigned_keys[i]); + } + } + + // The prefix (i.e., the most significant bits) of the k-th key + bit_ordered_type kth_prefix{}; + // The prefix mask (i.e., the bit mask with the most significant bits populated) of the k-th key + bit_ordered_type prefix_mask{}; + // The total number of items that compare strictly less than the k-th key's prefix (i.e., the number of items that + // are guaranteed to be selected) + int total_selected{}; + // The number of candidates that compare equal to the k-th key's prefix + auto num_candidates = IsFullTile ? tile_items : valid_items; + + // Identify the prefix of the k-th key + get_kth_key_prefix( + unsigned_keys, + k, + valid_items, + begin_bit, + end_bit, + total_selected, + num_candidates, + kth_prefix, + prefix_mask, + decomposer); + + // Scatter indices of selected items into shared memory (only for selecting key-value pairs, using a two-phase + // approach to lower shared memory requirements). + [[maybe_unused]] int scatter_indices[items_per_thread]; + if constexpr (!keys_only) + { + for (int i = 0; i < items_per_thread; ++i) + { + scatter_indices[i] = -1; + } + } + + // If all candidates are amongst the remaining top-k, we can simply select all items that compare less than or equal + // to the splitter prefix. Otherwise, we have to make sure that *all* candidates that compare strictly less than the + // splitter prefix are selected, and then select amongst candidates that compare equal to the splitter prefix to + // fill up the remaining slots up to k. + const bool select_all_candidates = expand_k_to_include_ties || num_candidates + total_selected == k; + + if (linear_tid == 0) + { + // Write offsets for selected items with key_prefix < kth_prefix + storage.stage.select.selected_offset[0] = 0; + // Write offsets for tied items across the k-th position, i.e., key_prefix == kth_prefix + storage.stage.select.selected_offset[1] = total_selected; + } + // Ensure atomic selection counter has been reset + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items_per_thread; ++i) + { + const bit_ordered_type key_prefix = unsigned_keys[i] & prefix_mask; + + const bool is_valid = (IsFullTile || linear_tid * items_per_thread + i < valid_items); + using comparison_t = ::cuda::std:: + conditional_t, ::cuda::std::greater<>>; + const bool is_selected = comparison_t{}(key_prefix, kth_prefix); + const bool is_candidate = key_prefix == kth_prefix; + + // We differentiate between candidates and selected only if not all candidates make it into the top-k items. + int item_class = (!select_all_candidates) && is_candidate ? 1 : 0; + + // Untwiddle the key before storing in shared memory + unsigned_keys[i] = bit_ordered_conversion::from_bit_ordered(decomposer, unsigned_keys[i]); + + if (is_valid && (is_selected || is_candidate)) + { + const histo_counter_t selected_offset = atomicAdd(&storage.stage.select.selected_offset[item_class], 1); + if constexpr (::cuda::is_floating_point_v) + { + storage.stage.select.exchange.keys[selected_offset] = + (flip_back_bits[i / 32] & (1u << (i % 32))) ? KeyT(-0.0) : ::cuda::std::bit_cast(unsigned_keys[i]); + } + else + { + storage.stage.select.exchange.keys[selected_offset] = ::cuda::std::bit_cast(unsigned_keys[i]); + } + if constexpr (!keys_only) + { + scatter_indices[i] = selected_offset; + } + } + } + + // Ensure all threads have finished writing to shared memory + __syncthreads(); + + // Gather selected items into thread registers for return. + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items_per_thread; ++i) + { + const int buffer_idx = linear_tid * items_per_thread + i; + if (buffer_idx < k) + { + keys[i] = storage.stage.select.exchange.keys[buffer_idx]; + } + } + + if constexpr (!keys_only) + { + // Ensure all keys have been loaded from shared memory before we repurpose the exchange buffer for values + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items_per_thread; ++i) + { + if (scatter_indices[i] >= 0) + { + storage.stage.select.exchange.values[scatter_indices[i]] = values[i]; + } + } + + // Ensure all values have been written to shared memory before we read them back in + __syncthreads(); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < items_per_thread; ++i) + { + const int buffer_idx = linear_tid * items_per_thread + i; + if (buffer_idx < k) + { + values[i] = storage.stage.select.exchange.values[buffer_idx]; + } + } + } + } + +public: + struct TempStorage : Uninitialized + {}; + + _CCCL_DEVICE_API _CCCL_FORCEINLINE block_topk_air(TempStorage& storage) + : storage(storage.Alias()) + , linear_tid(RowMajorTid(ThreadsPerBlock, 1, 1)) + {} + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void + select_keys(KeyT (&keys)[items_per_thread], int k, int valid_items, int begin_bit = 0, int end_bit = sizeof(KeyT) * 8) + { + NullType values[ItemsPerThread]; + select_topk(keys, values, k, valid_items, begin_bit, end_bit); + } + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void select_pairs( + KeyT (&keys)[items_per_thread], + ValueT (&values)[items_per_thread], + int k, + int valid_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8) + { + select_topk(keys, values, k, valid_items, begin_bit, end_bit); + } +}; +} // namespace detail +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/cub.cuh b/qwen3_6_scripts/cccl_preload/include/cub/cub.cuh new file mode 100644 index 00000000..92a6f018 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/cub.cuh @@ -0,0 +1,98 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * \file + * CUB umbrella include file + */ + +#pragma once + +// Static configuration +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include the specific device header instead (e.g. ). 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 + +// Block +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +// #include + +// Device +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Grid +#include +#include +#include + +// Thread +#include +#include +#include +#include +#include + +// Warp +#include +#include +#include +#include +#include +#include + +// Iterator +#include +#include +#include +#include + +// Util +#include +#include +#include +#include +#include +#include diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/array_utils.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/array_utils.cuh new file mode 100644 index 00000000..98683773 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/array_utils.cuh @@ -0,0 +1,54 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include // static_size_v +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + +/*********************************************************************************************************************** + * Generic Array-like to Array Conversion + **********************************************************************************************************************/ + +template +[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::array> +to_array_impl(const Input& input, ::cuda::std::index_sequence) +{ + using ArrayType = ::cuda::std::array>; + return ArrayType{static_cast(input[i])...}; +} + +template +[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::array> +to_array(const Input& input) +{ + using InputType = ::cuda::std::iter_value_t; + using CastType1 = ::cuda::std::_If<::cuda::std::is_same_v, InputType, CastType>; + return to_array_impl(input, ::cuda::std::make_index_sequence>{}); +} + +#endif // !_CCCL_DOXYGEN_INVOKED +} // namespace detail +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/binary_search_helpers.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/binary_search_helpers.cuh new file mode 100644 index 00000000..fdf7f1dc --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/binary_search_helpers.cuh @@ -0,0 +1,126 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::find +{ +constexpr ::cuda::std::ptrdiff_t linear_lower_bound_threshold = 8; + +template +struct comp_wrapper_t +{ + RangeIteratorT first; + RangeNumItemsT num_items; + CompareOpT op; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(::cuda::std::tuple args) const + { + using DifferenceT = ::cuda::std::iter_difference_t; + const auto last = first + static_cast(num_items); + + ::cuda::std::get<1>(args) = Mode::Invoke(first, last, ::cuda::std::get<0>(args), op); + } +}; + +template +_CCCL_HOST_DEVICE auto make_comp_wrapper(RangeIteratorT first, RangeNumItemsT num_items, CompareOpT comp) +{ + return comp_wrapper_t{first, num_items, comp}; +} + +struct lower_bound +{ + template + _CCCL_DEVICE _CCCL_FORCEINLINE static DifferenceT + Linear(RangeIteratorT first, DifferenceT num_items, const T& value, CompareOpT comp) + { + DifferenceT retval = 0; + for (DifferenceT i = 0; i < num_items; ++i) + { + retval += static_cast(comp(first[i], value)); + } + + return retval; + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE static ::cuda::std::ptrdiff_t + Invoke(RangeIteratorT first, RangeIteratorT last, const T& value, CompareOpT comp) + { + return ::cuda::std::lower_bound(first, last, value, comp) - first; + } +}; + +struct upper_bound +{ + template + _CCCL_DEVICE _CCCL_FORCEINLINE static DifferenceT + Linear(RangeIteratorT first, DifferenceT num_items, const T& value, CompareOpT comp) + { + DifferenceT retval = 0; + for (DifferenceT i = 0; i < num_items; ++i) + { + retval += static_cast(!comp(value, first[i])); + } + + return retval; + } + + template + _CCCL_DEVICE _CCCL_FORCEINLINE static ::cuda::std::ptrdiff_t + Invoke(RangeIteratorT first, RangeIteratorT last, const T& value, CompareOpT comp) + { + return ::cuda::std::upper_bound(first, last, value, comp) - first; + } +}; + +template +struct binary_search_transform_op_t +{ + RangeIteratorT first; + RangeNumItemsT num_items; + CompareOpT op; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE ::cuda::std::ptrdiff_t operator()(const Value& value) const + { + using DifferenceT = ::cuda::std::iter_difference_t; + const auto count = static_cast(num_items); + + if (num_items <= static_cast(linear_lower_bound_threshold)) + { + return Mode::Linear(first, count, value, op); + } + + return Mode::Invoke(first, first + count, value, op); + } +}; + +template +_CCCL_HOST_DEVICE auto make_binary_search_transform_op(RangeIteratorT first, RangeNumItemsT num_items, CompareOpT comp) +{ + return binary_search_transform_op_t{first, num_items, comp}; +} +} // namespace detail::find + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/cc_dispatch.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/cc_dispatch.cuh new file mode 100644 index 00000000..37e385cf --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/cc_dispatch.cuh @@ -0,0 +1,160 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +// makes a functor that gets the policy for CC from PolicySelector when called +template +struct policy_getter : PolicySelector +{ + _CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE constexpr auto operator()() const + { + return PolicySelector::operator()(::cuda::compute_capability{CC}); + } +}; + +// Device-only variant for kernel-side compile-time policy queries. +template +struct device_policy_getter : PolicySelector +{ + _CCCL_DEVICE_API _CCCL_FORCEINLINE constexpr auto operator()() const + { + return PolicySelector::operator()(::cuda::compute_capability{CC}); + } +}; + +#if !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) +# if _CCCL_STD_VER < 2020 && !_CCCL_COMPILER(GCC, <, 8) +template +struct lowest_cc_resolver; + +// we keep the compile-time build up of the mapping table outside a template parameterized by a user-provided callable +template +struct lowest_cc_resolver<::cuda::std::integer_sequence, PolicySelector, Is...> +{ + static_assert(sizeof...(CudaCcs) == sizeof...(Is)); + + using policy_t = decltype(PolicySelector{}(::cuda::compute_capability{})); + + static constexpr ::cuda::compute_capability all_ccs[sizeof...(Is)]{::cuda::compute_capability{CudaCcs}...}; + static constexpr policy_t all_policies[sizeof...(Is)]{PolicySelector{}(all_ccs[Is])...}; + + _CCCL_HOST_DEVICE_API static constexpr auto find_lowest(size_t i) -> ::cuda::compute_capability + { + const auto& policy = all_policies[i]; + while (i > 0 && policy == all_policies[i - 1]) + { + --i; + } + return all_ccs[i]; + } + + static constexpr ::cuda::compute_capability lowest_cc_with_same_policy[sizeof...(Is)]{find_lowest(Is)...}; +}; +# endif // if _CCCL_STD_VER < 2020 && !_CCCL_COMPILER(GCC, <, 8) + +// GCC below 12 ICEs in some cases when creating an integral_constant holding a policy +# if _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12) +template +struct policy_constant +{ + _CCCL_API constexpr auto operator()() const noexcept + { + return P; + } +}; +# else // _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12) +template // using will miscompile on GCC 12 +using policy_constant = ::cuda::std::integral_constant; +# endif // _CCCL_STD_VER >= 2020 && _CCCL_COMPILER(GCC, <, 12) + +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_to_cc_list( + PolicySelector policy_selector, ::cuda::compute_capability device_cc, FunctorT&& f, ::cuda::std::index_sequence) +{ + constexpr auto all_ccs = ::cuda::__target_compute_capabilities(); + + _CCCL_ASSERT(((device_cc == all_ccs[Is]) || ...), + "device_cc must appear in the list of compute capabilities compiled for"); + + cudaError_t e = cudaErrorInvalidDeviceFunction; +# if _CCCL_STD_VER >= 2020 + // In C++20, we just create an integral_constant holding the policy, because policies are structural types in C++20. + // This causes f to be only instantiated for each distinct policy, since the same policy for different arches results + // in the same integral_constant type passed to f + using policy_t = decltype(policy_selector(::cuda::compute_capability{})); + (..., (device_cc == all_ccs[Is] ? (e = f(policy_constant{})) : cudaSuccess)); +# else // _CCCL_STD_VER >= 2020 +# if _CCCL_COMPILER(GCC, <, 8) + // GCC 7 ICEs on constexpr evaluation of policy comparisons, so we skip the lowest-CC-with-same-policy optimization + // and instantiate f for each CC directly. This may increase compile time and binary size. + (..., + (device_cc == all_ccs[Is] ? (e = f(policy_getter{policy_selector})) + : cudaSuccess)); +# else // _CCCL_COMPILER(GCC, <, 8) + // In C++17, we have to collapse architectures with the same policies ourselves, so we instantiate call_for_cc once + // per policy on the lowest CC which produces the same policy + using resolver_t = + lowest_cc_resolver<::cuda::std::integer_sequence, PolicySelector, Is...>; + (..., + (device_cc == all_ccs[Is] + ? (e = f(policy_getter{policy_selector})) + : cudaSuccess)); +# endif // _CCCL_COMPILER(GCC, <, 8) +# endif // _CCCL_STD_VER >= 2020 + return e; +} + +//! Takes a policy hub and instantiates f with the minimum possible number of nullary functor types that return a policy +//! at compile-time (if possible), and then calls the appropriate instantiation based on a runtime GPU architecture. +//! Depending on the used compiler, C++ standard, and available macros, a different number of instantiations may be +//! produced. +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t +dispatch_compute_cap(PolicySelector policy_selector, ::cuda::compute_capability device_cc, F&& f) +{ + // when not using CCCL.C, policy_selector is empty since all information is contained in its type + static_assert(::cuda::std::is_empty_v); + return dispatch_to_cc_list( + policy_selector, + device_cc, + ::cuda::std::forward(f), + ::cuda::std::make_index_sequence<::cuda::__target_compute_capabilities().size()>{}); +} + +#else // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) + +// if we are compiling CCCL.C with runtime policies, we cannot query the policy hub at compile time +_CCCL_EXEC_CHECK_DISABLE +template +_CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE cudaError_t +dispatch_compute_cap(PolicySelector policy_selector, ::cuda::compute_capability device_cc, F&& f) +{ + return f([&] { + return policy_selector(device_cc); + }); +} +#endif // !defined(CUB_DEFINE_RUNTIME_POLICIES) && !_CCCL_COMPILER(NVRTC) +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/choose_offset.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/choose_offset.cuh new file mode 100644 index 00000000..c029f19c --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/choose_offset.cuh @@ -0,0 +1,135 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +/** + * choose_offset checks NumItemsT, the type of the num_items parameter, and + * selects the offset type based on it. + */ +template +struct choose_offset +{ + // NumItemsT must be an integral type (but not bool). + static_assert(::cuda::std::is_integral_v + && !::cuda::std::is_same_v<::cuda::std::remove_cv_t, bool>, + "NumItemsT must be an integral type, but not bool"); + + // Unsigned integer type for global offsets. + using type = ::cuda::std::_If<(sizeof(NumItemsT) <= 4), uint32_t, unsigned long long>; +}; + +/** + * choose_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and + * selects the offset type based on it. + */ +template +using choose_offset_t = typename choose_offset::type; + +/** + * promote_small_offset checks NumItemsT, the type of the num_items parameter, and + * promotes any integral type smaller than 32 bits to a signed 32-bit integer type. + */ +template +struct promote_small_offset +{ + // NumItemsT must be an integral type (but not bool). + static_assert(::cuda::std::is_integral_v + && !::cuda::std::is_same_v<::cuda::std::remove_cv_t, bool>, + "NumItemsT must be an integral type, but not bool"); + + // Unsigned integer type for global offsets. + using type = ::cuda::std::_If<(sizeof(NumItemsT) < 4), int32_t, NumItemsT>; +}; + +/** + * promote_small_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and + * promotes any integral type smaller than 32 bits to a signed 32-bit integer type. + */ +template +using promote_small_offset_t = typename promote_small_offset::type; + +/** + * choose_signed_offset checks NumItemsT, the type of the num_items parameter, and + * selects the offset type to be either int32 or int64, such that the selected offset type covers the range of NumItemsT + * unless it was uint64, in which case int64 will be used. + */ +template +struct choose_signed_offset +{ + // NumItemsT must be an integral type (but not bool). + static_assert(::cuda::std::is_integral_v + && !::cuda::std::is_same_v<::cuda::std::remove_cv_t, bool>, + "NumItemsT must be an integral type, but not bool"); + + // Signed integer type for global offsets. + // uint32 -> int64, else + // LEQ 4B -> int32, else + // int64 + using type = ::cuda::std::_If<(::cuda::std::is_integral_v && ::cuda::std::is_unsigned_v), + ::cuda::std::int64_t, + ::cuda::std::_If<(sizeof(NumItemsT) <= 4), ::cuda::std::int32_t, ::cuda::std::int64_t>>; + + /** + * Checks if the given num_items can be covered by the selected offset type. If not, returns cudaErrorInvalidValue, + * otherwise returns cudaSuccess. + */ + static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE cudaError_t is_exceeding_offset_type(NumItemsT num_items) + { + _CCCL_DIAG_PUSH + _CCCL_DIAG_SUPPRESS_MSVC(4127) /* conditional expression is constant */ + if (sizeof(NumItemsT) >= 8 && num_items > static_cast(::cuda::std::numeric_limits::max())) + { + return cudaErrorInvalidValue; + } + _CCCL_DIAG_POP + return cudaSuccess; + } +}; + +/** + * choose_signed_offset_t is an alias template that checks NumItemsT, the type of the num_items parameter, and + * selects the corresponding signed offset type based on it. + */ +template +using choose_signed_offset_t = typename choose_signed_offset::type; + +/** + * common_iterator_value sets member type to the common_type of + * value_type for all argument types. used to get OffsetT in + * DeviceSegmentedReduce. + */ +template +struct common_iterator_value +{ + using type = ::cuda::std::common_type_t<::cuda::std::__iter_value_type...>; +}; +template +using common_iterator_value_t = typename common_iterator_value::type; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/deferred_parameter.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/deferred_parameter.cuh new file mode 100644 index 00000000..e808cbd2 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/deferred_parameter.cuh @@ -0,0 +1,82 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +#if !_CCCL_COMPILER(NVRTC) +// Preserve deferred problem sizes for dispatch and canonicalize immediate values to CUB's offset type. +template +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto make_num_items_dispatch_arg(NumItemsT num_items) noexcept +{ + using args_traits_t = ::cuda::args::__traits; + + if constexpr (args_traits_t::is_deferred) + { + return num_items; + } + else + { + using offset_t = choose_offset_t; + return static_cast(::cuda::args::__unwrap(num_items)); + } +} + +// Forms a kernel parameter from a single-value argument without reading a deferred source. +// Immediate values are converted to TargetT. Deferred arguments are unwrapped to their source, erasing bounds from +// the kernel type and payload. +template +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE constexpr auto parameter_from_host(ParameterT parameter) noexcept +{ + using args_traits_t = ::cuda::args::__traits; + static_assert(args_traits_t::is_single_value, "parameter must contain a single value"); + + if constexpr (args_traits_t::is_deferred) + { + return ::cuda::args::__unwrap(parameter); + } + else + { + return static_cast(::cuda::args::__unwrap(parameter)); + } +} + +template +using parameter_from_host_t = decltype(parameter_from_host(::cuda::std::declval())); +#endif // !_CCCL_COMPILER(NVRTC) + +// Forms a value from a kernel parameter, reading element zero when the parameter is a deferred source. +template +[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE TargetT parameter_from_device(ParameterT parameter) noexcept +{ + if constexpr (::cuda::std::is_same_v) + { + return parameter; + } + else + { + return static_cast(parameter[0]); + } +} +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/delay_constructor.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/delay_constructor.cuh new file mode 100644 index 00000000..721c09d8 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/delay_constructor.cuh @@ -0,0 +1,150 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +template +inline constexpr auto lookback_delay_policy_from_type = 0; + +template +inline constexpr auto lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::no_delay, 0, L2WriteLatency}; + +template +inline constexpr auto lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::fixed_delay, Delay, L2WriteLatency}; + +template +inline constexpr auto lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff, Delay, L2WriteLatency}; + +template +inline constexpr auto lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff_jitter, Delay, L2WriteLatency}; + +template +inline constexpr auto + lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backoff_jitter_window, Delay, L2WriteLatency}; + +template +inline constexpr auto + lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon_jitter_window, Delay, L2WriteLatency}; + +template +inline constexpr auto lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon_jitter, Delay, L2WriteLatency}; + +template +inline constexpr auto lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::exponential_backon, Delay, L2WriteLatency}; + +template +inline constexpr auto + lookback_delay_policy_from_type> = + LookbackDelayPolicy{LookbackDelayAlgorithm::__reduce_by_key, Delay, L2WriteLatency}; + +template +struct delay_constructor_for; + +template +struct delay_constructor_for +{ + using type = no_delay_constructor_t; +}; + +template +struct delay_constructor_for +{ + using type = fixed_delay_constructor_t; +}; + +template +struct delay_constructor_for +{ + using type = exponential_backoff_constructor_t; +}; + +template +struct delay_constructor_for +{ + using type = exponential_backoff_jitter_constructor_t; +}; + +template +struct delay_constructor_for +{ + using type = exponential_backoff_jitter_window_constructor_t; +}; + +template +struct delay_constructor_for +{ + using type = exponential_backon_jitter_window_constructor_t; +}; + +template +struct delay_constructor_for +{ + using type = exponential_backon_jitter_constructor_t; +}; + +template +struct delay_constructor_for +{ + using type = exponential_backon_constructor_t; +}; + +template +struct delay_constructor_for +{ + using type = reduce_by_key_delay_constructor_t; +}; + +template +using delay_constructor_t = typename delay_constructor_for::type; + +_CCCL_HOST_DEVICE_API constexpr auto default_delay_constructor_policy(bool is_primitive_or_trivially_copyable) +{ + if (is_primitive_or_trivially_copyable) + { + return LookbackDelayPolicy{LookbackDelayAlgorithm::fixed_delay, 350, 450}; + } + return LookbackDelayPolicy{LookbackDelayAlgorithm::no_delay, 0, 450}; +} + +_CCCL_HOST_DEVICE_API constexpr auto default_reduce_by_key_delay_constructor_policy( + int key_size, + int value_size, + bool key_is_primitive_or_trivially_copyable, + bool value_is_primitive_or_trivially_copyable) +{ + if (value_is_primitive_or_trivially_copyable && (value_size + key_size < 16)) + { + return LookbackDelayPolicy{LookbackDelayAlgorithm::__reduce_by_key, 350, 450}; + } + return default_delay_constructor_policy(key_is_primitive_or_trivially_copyable); +} +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/device_double_buffer.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/device_double_buffer.cuh new file mode 100644 index 00000000..f2b1ccd3 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/device_double_buffer.cuh @@ -0,0 +1,94 @@ +/* + * Copyright 2021 NVIDIA Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +/** + * @brief It's a double-buffer storage wrapper for multi-pass stream + * transformations that require more than one storage array for + * streaming intermediate results back and forth. + * + * Many multi-pass computations require a pair of "ping-pong" storage buffers + * (e.g., one for reading from and the other for writing to, and then + * vice-versa for the subsequent pass). This structure wraps a set of device + * buffers. + * + * Unlike `cub::DoubleBuffer` this class doesn't provide a "selector" member + * to track which buffer is "current". The main reason for this class existence + * is the performance difference. Since `cub::DoubleBuffer` relies on the + * runtime variable to index pointers arrays, they are placed in the local + * memory instead of registers. Local memory accesses significantly affect + * performance. On the contrary, this class swaps pointer, so all operations + * can be performed in registers. + */ +template +class device_double_buffer +{ + /// Pair of device buffer pointers + T* m_current_buffer{}; + T* m_alternate_buffer{}; + +public: + /** + * @param d_current + * The currently valid buffer + * + * @param d_alternate + * Alternate storage buffer of the same size as @p d_current + */ + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE device_double_buffer(T* current, T* alternate) + : m_current_buffer(current) + , m_alternate_buffer(alternate) + {} + + /// \brief Return pointer to the currently valid buffer + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE T* current() const + { + return m_current_buffer; + } + + /// \brief Return pointer to the currently invalid buffer + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE T* alternate() const + { + return m_alternate_buffer; + } + + _CCCL_HOST_DEVICE void swap() + { + T* tmp = m_current_buffer; + m_current_buffer = m_alternate_buffer; + m_alternate_buffer = tmp; + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/device_memory_resource.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/device_memory_resource.cuh new file mode 100644 index 00000000..af584341 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/device_memory_resource.cuh @@ -0,0 +1,97 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +// TODO(gevtushenko/srinivasyadav18): move cudax `device_memory_resource` to `cuda::__device_memory_resource` and remove +// this implementation +struct device_memory_resource +{ + CUB_RUNTIME_FUNCTION void* allocate(size_t bytes, size_t /* alignment */) + { + void* ptr{nullptr}; + NV_IF_ELSE_TARGET( + NV_IS_HOST, + (_CCCL_TRY_CUDA_API(::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, NULL);), + ({ + _CubLog("%s\n", "cub::detail::device_memory_resource::allocate not supported from device code."); + ::cuda::std::terminate(); + })); + _CCCL_ASSERT(ptr != nullptr, "allocate failed to allocate with cudaMallocAsync"); + return ptr; + } + + CUB_RUNTIME_FUNCTION void deallocate(void* ptr, size_t /* bytes */) + { + NV_IF_ELSE_TARGET( // + NV_IS_HOST, + (_CCCL_TRY_CUDA_API(::cudaFree, "deallocate failed", ptr);), + ({ + _CubLog("%s\n", "cub::detail::device_memory_resource::deallocate not supported from device code."); + ::cuda::std::terminate(); + })); + } + + CUB_RUNTIME_FUNCTION void* allocate(::cuda::stream_ref stream, size_t bytes, size_t /* alignment */) + { + return allocate(stream, bytes); + } + + CUB_RUNTIME_FUNCTION void* allocate(::cuda::stream_ref stream, size_t bytes) + { + void* ptr{nullptr}; + NV_IF_ELSE_TARGET( // + NV_IS_HOST, + ({ + _CCCL_TRY_CUDA_API( + ::cudaMallocAsync, "allocate failed to allocate with cudaMallocAsync", &ptr, bytes, stream.get()); + }), + ({ + _CubLog("%s\n", "cub::detail::device_memory_resource::allocate not supported from device code."); + ::cuda::std::terminate(); + })); + return ptr; + } + + CUB_RUNTIME_FUNCTION void deallocate(::cuda::stream_ref stream, void* ptr, size_t bytes, size_t /* alignment */) + { + deallocate(stream, ptr, bytes); + } + + CUB_RUNTIME_FUNCTION void deallocate(::cuda::stream_ref stream, void* ptr, size_t /* bytes */) + { + NV_IF_ELSE_TARGET( // + NV_IS_HOST, + (_CCCL_TRY_CUDA_API(::cudaFreeAsync, "deallocate failed", ptr, stream.get());), + ({ + _CubLog("%s\n", "cub::detail::device_memory_resource::deallocate not supported from device code."); + ::cuda::std::terminate(); + })); + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/env_dispatch.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/env_dispatch.cuh new file mode 100644 index 00000000..1862a4ff --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/env_dispatch.cuh @@ -0,0 +1,137 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +//! @cond +//! Generic environment-based algorithm dispatch wrapper +//! +//! Handles common boilerplate for all env-based algorithms: +//! - Query stream, memory resource, and tuning from environment +//! - Two-phase call (query temp storage size, then execute) +//! - Temporary storage allocation/deallocation +//! - Memory resource querying from environment +//! +//! @param env The execution environment +//! @param algorithm_callable Callable that invokes the algorithm implementation with determinism specified +template +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env(const EnvT& env, AlgorithmCallable&& algorithm_callable) +{ + // Query stream from environment + auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); + + // Query memory resource from environment + auto mr = ::cuda::__call_or(::cuda::mr::__get_memory_resource, detail::device_memory_resource{}, env); + + // Query tuning from environment + const auto tuning = ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env); + + void* d_temp_storage = nullptr; + size_t temp_storage_bytes = 0; + + // Phase 1: Query temporary storage size + if (const auto error = algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get())) + { + return error; + } + + // Allocate temporary storage + if (const auto error = CubDebug(detail::temporary_storage::allocate(stream, d_temp_storage, temp_storage_bytes, mr))) + { + return error; + } + + // Phase 2: Execute algorithm + const auto error = algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get()); + + // Deallocate temporary storage (always attempt, even on error) + const auto deallocate_error = + CubDebug(detail::temporary_storage::deallocate(stream, d_temp_storage, temp_storage_bytes, mr)); + + // Algorithm error takes precedence over deallocation error + return (error != cudaSuccess) ? error : deallocate_error; +} +//! @endcond + +template +CUB_RUNTIME_FUNCTION static cudaError_t +dispatch_with_env_and_tuning(const EnvT& env, AlgorithmCallable&& algorithm_callable) +{ + return detail::dispatch_with_env( + env, + [&algorithm_callable]( + [[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + using policy_t = decltype(DefaultPolicySelector{}(::cuda::compute_capability{})); + using policy_selector = + ::cuda::std::execution::__query_result_or_t; + return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream); + }); +} + +//! @cond +//! Generic environment-based algorithm dispatch wrapper +//! +//! Handles common boilerplate for env-based algorithms with user provided memory: +//! - Query stream, and tuning from environment +//! - Single-phase call passing user provided memory and size +//! +//! @param env The execution environment +//! @param[in] d_temp_storage @devicestorage +//! @param[in,out] temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation +//! @param algorithm_callable Callable that invokes the algorithm implementation with determinism specified +template +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env( + void* d_temp_storage, size_t& temp_storage_bytes, const EnvT& env, AlgorithmCallable&& algorithm_callable) +{ + // Query stream from environment + auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); + + // Query tuning from environment + const auto tuning = ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env); + + return algorithm_callable(tuning, d_temp_storage, temp_storage_bytes, stream.get()); +} +//! @endcond + +template +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_with_env_and_tuning( + void* d_temp_storage, size_t& temp_storage_bytes, const EnvT& env, AlgorithmCallable&& algorithm_callable) +{ + return detail::dispatch_with_env( + d_temp_storage, + temp_storage_bytes, + env, + [&algorithm_callable]( + [[maybe_unused]] auto tuning_env, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + using policy_t = decltype(DefaultPolicySelector{}(::cuda::compute_capability{})); + using policy_selector = + ::cuda::std::execution::__query_result_or_t; + return algorithm_callable(policy_selector{}, d_temp_storage, temp_storage_bytes, stream); + }); +} +//! @endcond +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/fast_modulo_division.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/fast_modulo_division.cuh new file mode 100644 index 00000000..8b26b0bb --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/fast_modulo_division.cuh @@ -0,0 +1,226 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include // implicit_prom_t +#include // _CCCL_HAS_INT128() + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include // CHAR_BIT +#include // uint64_t +#include + +#if defined(CCCL_ENABLE_DEVICE_ASSERTIONS) +_CCCL_BEGIN_NV_DIAG_SUPPRESS(186) // pointless comparison of unsigned integer with zero +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +/*********************************************************************************************************************** + * larger_unsigned_type + **********************************************************************************************************************/ + +template +struct larger_unsigned_type +{ + using type = void; +}; + +template +struct larger_unsigned_type> +{ + using type = ::cuda::std::uint32_t; +}; + +template +struct larger_unsigned_type> +{ + using type = ::cuda::std::uint64_t; +}; + +#if _CCCL_HAS_INT128() + +template +struct larger_unsigned_type> +{ + using type = __uint128_t; +}; + +#endif // _CCCL_HAS_INT128() + +template +using larger_unsigned_type_t = typename larger_unsigned_type::type; + +template +using unsigned_implicit_prom_t = ::cuda::std::make_unsigned_t>; + +template +using supported_integral = + ::cuda::std::bool_constant<::cuda::std::is_integral_v && !::cuda::std::is_same_v && (sizeof(T) <= 8)>; + +/*********************************************************************************************************************** + * Extract higher bits after multiplication + **********************************************************************************************************************/ + +template +[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE unsigned_implicit_prom_t +multiply_extract_higher_bits(T value, R multiplier) +{ + static_assert(supported_integral::value, "unsupported type"); + static_assert(supported_integral::value, "unsupported type"); + if constexpr (::cuda::std::is_signed_v) + { + _CCCL_ASSERT(value >= 0, "value must be non-negative"); + } + if constexpr (::cuda::std::is_signed_v) + { + _CCCL_ASSERT(multiplier >= 0, "multiplier must be non-negative"); + } + static constexpr int NumBits = sizeof(DivisorType) * CHAR_BIT; + using unsigned_t = unsigned_implicit_prom_t; + using larger_t = larger_unsigned_type_t; + // clang-format off + NV_IF_ELSE_TARGET( + NV_IS_HOST, + (return static_cast((static_cast(value) * multiplier) >> NumBits);), + ({return (sizeof(T) == 8) + ? static_cast(__umul64hi(value, multiplier)) + : static_cast((static_cast(value) * multiplier) >> NumBits);})); + // clang-format on +} + +/*********************************************************************************************************************** + * Fast Modulo/Division based on Precomputation + **********************************************************************************************************************/ + +_CCCL_DIAG_PUSH +_CCCL_DIAG_SUPPRESS_MSVC(4127) /* conditional expression is constant */ + +template +class fast_div_mod +{ + static_assert(supported_integral::value, "unsupported type"); + + // uint16_t is a special case that would requires complex logic. Workaround: convert to int + using T = ::cuda::std::conditional_t<::cuda::std::is_same_v, int, T1>; + using unsigned_t = unsigned_implicit_prom_t; + +public: + template + struct result + { + using common_t = decltype(R{} / T{}); + common_t quotient; + common_t remainder; + }; + + fast_div_mod() = delete; + + _CCCL_HOST_DEVICE explicit fast_div_mod(T divisor) noexcept + : _divisor{static_cast(divisor)} + { + using larger_t = larger_unsigned_type_t; + _CCCL_ASSERT(divisor > 0, "divisor must be positive"); + auto udivisor = static_cast(divisor); + // the following branches are needed to avoid negative shift + if (::cuda::is_power_of_two(udivisor)) + { + _shift_right = ::cuda::std::bit_width(udivisor) - 1; + return; + } + else if (sizeof(T) == 8 && divisor == 3) + { + return; + } + constexpr int BitSize = sizeof(T) * CHAR_BIT; // 32 + constexpr int BitOffset = BitSize / 16; // 2 + int num_bits = ::cuda::std::bit_width(udivisor) + 1; + _CCCL_ASSERT(static_cast(num_bits + BitSize - BitOffset) < sizeof(larger_t) * CHAR_BIT, "overflow error"); + // without explicit power-of-two check, num_bits needs to replace +1 with !::cuda::is_power_of_two(udivisor) + _multiplier = static_cast(::cuda::ceil_div(larger_t{1} << (num_bits + BitSize - BitOffset), // + static_cast(divisor))); + _shift_right = num_bits - BitOffset; + _CCCL_ASSERT(_multiplier != 0, "overflow error"); + } + + fast_div_mod(const fast_div_mod&) noexcept = default; + + fast_div_mod(fast_div_mod&&) noexcept = default; + + template + [[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE result operator()(R dividend) const noexcept + { + static_assert(supported_integral::value, "unsupported type"); + using common_t = decltype(R{} / T{}); + using ucommon_t = ::cuda::std::make_unsigned_t; + using result_t = result; + _CCCL_ASSERT(dividend >= 0, "divisor must be non-negative"); + auto udividend = static_cast(dividend); + if (_divisor == 1) + { + return result_t{static_cast(dividend), common_t{}}; + } + else if (_divisor > unsigned_t{::cuda::std::numeric_limits::max() / 2}) + { + auto quotient = udividend >= static_cast(_divisor); + return result_t{static_cast(quotient), static_cast(udividend - (quotient * _divisor))}; + } + else if (sizeof(T) == 8 && _divisor == 3) + { + return result_t{static_cast(udividend / 3), static_cast(udividend % 3)}; + } + auto higher_bits = (_multiplier == 0) ? udividend : multiply_extract_higher_bits(dividend, _multiplier); + auto quotient = higher_bits >> _shift_right; + auto remainder = udividend - (quotient * _divisor); + _CCCL_ASSERT(quotient == udividend / _divisor, "wrong quotient"); + _CCCL_ASSERT(remainder < (ucommon_t) _divisor, "remainder out of range"); + return result_t{static_cast(quotient), static_cast(remainder)}; + } + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE friend implicit_prom_t operator/(R dividend, fast_div_mod div) noexcept + { + return div(dividend).quotient; + } + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE friend implicit_prom_t operator%(R dividend, fast_div_mod div) noexcept + { + return div(dividend).remainder; + } + +private: + unsigned_t _divisor = 1; + unsigned_t _multiplier = 0; + unsigned _shift_right = 0; +}; +_CCCL_DIAG_POP +} // namespace detail + +CUB_NAMESPACE_END + +#if defined(CCCL_ENABLE_DEVICE_ASSERTIONS) +_CCCL_END_NV_DIAG_SUPPRESS() +#endif // CCCL_ENABLE_DEVICE_ASSERTIONS diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/launcher/cuda_runtime.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/launcher/cuda_runtime.cuh new file mode 100644 index 00000000..ab8e3a95 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/launcher/cuda_runtime.cuh @@ -0,0 +1,147 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +struct TripleChevronFactory +{ + CUB_RUNTIME_FUNCTION void __assert_pdl_allowed(bool dependent_launch) const + { + if (dependent_launch) + { + [[maybe_unused]] int sm_version = 0; + _CCCL_ASSERT(SmVersion(sm_version) == cudaSuccess, "Failed to query SM compute capability"); + if (sm_version >= 900) + { + [[maybe_unused]] ::cuda::compute_capability cc; + _CCCL_ASSERT(PtxComputeCap(cc) == cudaSuccess, "Failed to query PTX compute capability"); + _CCCL_ASSERT((cc >= ::cuda::compute_capability{9, 0}), + "Enabling PDL for a kernel launch requires CC 9.0+ PTX/SASS when running on SM90+"); + } + } + } + + CUB_RUNTIME_FUNCTION THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron operator()( + dim3 grid, dim3 block, ::cuda::std::size_t shared_mem, ::cudaStream_t stream, bool dependent_launch = false) const + { + __assert_pdl_allowed(dependent_launch); + return THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid, block, shared_mem, stream, dependent_launch); + } + + template + CUB_RUNTIME_FUNCTION ::cudaError_t PtxVersion(int& version) + { + return cub::PtxVersion(version); + } + + template + CUB_RUNTIME_FUNCTION ::cudaError_t PtxComputeCap(::cuda::compute_capability& cc) const + { + return ptx_compute_cap(cc); + } + + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MultiProcessorCount(int& sm_count) const + { + int device_ordinal; + ::cudaError_t error = CubDebug(::cudaGetDevice(&device_ordinal)); + if (::cudaSuccess != error) + { + return error; + } + + // Get SM count + return ::cudaDeviceGetAttribute(&sm_count, ::cudaDevAttrMultiProcessorCount, device_ordinal); + } + + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t + MaxSmOccupancy(int& sm_occupancy, Kernel kernel_ptr, int block_size, int dynamic_smem_bytes = 0) + { + return ::cudaOccupancyMaxActiveBlocksPerMultiprocessor(&sm_occupancy, kernel_ptr, block_size, dynamic_smem_bytes); + } + + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t MaxGridDimX(int& max_grid_dim_x) const + { + int device_ordinal; + ::cudaError_t error = CubDebug(::cudaGetDevice(&device_ordinal)); + if (::cudaSuccess != error) + { + return error; + } + + // Get max grid dimension + return ::cudaDeviceGetAttribute(&max_grid_dim_x, ::cudaDevAttrMaxGridDimX, device_ordinal); + } + + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t + MemsetAsync(void* dst, unsigned char value, size_t num_bytes, ::cudaStream_t stream) const + { + return ::cudaMemsetAsync(dst, value, num_bytes, stream); + } + + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t + MemcpyAsync(void* dst, const void* src, size_t num_bytes, ::cudaMemcpyKind kind, ::cudaStream_t stream) const + { + return ::cudaMemcpyAsync(dst, src, num_bytes, kind, stream); + } + + // TODO(bgruber): this is very similar to thrust::cuda_cub::core::get_max_shared_memory_per_block. We should unify + // this. + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION cudaError_t MaxSharedMemory(int& max_shared_memory) const + { + int device = 0; + auto error = CubDebug(cudaGetDevice(&device)); + if (error != cudaSuccess) + { + return error; + } + + return cudaDeviceGetAttribute(&max_shared_memory, cudaDevAttrMaxSharedMemoryPerBlock, device); + } + + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t + max_dynamic_smem_size_for(int& max_dynamic_smem_size, [[maybe_unused]] Kernel kernel_ptr) + { + NV_IF_ELSE_TARGET(NV_IS_HOST, // + ({ return MaxPotentialDynamicSmemBytes(max_dynamic_smem_size, kernel_ptr); }), + ({ + ::cudaFuncAttributes func_attrs{}; + if (const auto error = CubDebug(::cudaFuncGetAttributes(&func_attrs, kernel_ptr))) + { + return error; + } + max_dynamic_smem_size = func_attrs.maxDynamicSharedSizeBytes; + return cudaSuccess; + })) + } + + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION ::cudaError_t set_max_dynamic_smem_size_for(Kernel kernel_ptr, int smem_size) + { + return CubDebug(::cudaFuncSetAttribute(kernel_ptr, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); + } +}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/mdspan_utils.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/mdspan_utils.cuh new file mode 100644 index 00000000..c6e95c43 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/mdspan_utils.cuh @@ -0,0 +1,113 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include // fast_div_mod + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +_CCCL_DIAG_PUSH +_CCCL_DIAG_SUPPRESS_MSVC(4702) // unreachable code (even if there are no branches!) + +// Compute the submdspan size of a given rank +template +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::make_unsigned_t +size_range(const ::cuda::std::extents& ext, int start, int end) +{ + _CCCL_ASSERT(start >= 0 && end <= static_cast(ext.rank()), "invalid start or end"); + ::cuda::std::make_unsigned_t s = 1; + for (auto i = start; i < end; i++) + { + s *= ext.extent(i); + } + return s; +} + +_CCCL_DIAG_POP // MSVC(4702) + + template + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::make_unsigned_t + size(const ::cuda::std::extents& ext) +{ + return cub::detail::size_range(ext, 0, static_cast(ext.rank())); +} + +template +[[nodiscard]] _CCCL_HOST_DEVICE_API auto sub_size_fast_div_mod_impl(const ::cuda::std::extents& ext) +{ + using fast_mod_div_t = fast_div_mod; + constexpr auto start = IsLayoutRight ? Position + 1 : 0; + constexpr auto end = IsLayoutRight ? sizeof...(E) : Position; + return fast_mod_div_t(cub::detail::size_range(ext, start, end)); +} + +// precompute modulo/division for each submdspan size (by rank) +template +[[nodiscard]] _CCCL_HOST_DEVICE_API auto +sub_sizes_fast_div_mod(const ::cuda::std::extents& ext, ::cuda::std::index_sequence = {}) +{ + using fast_mod_div_t = fast_div_mod; + using array_t = ::cuda::std::array; + return array_t{cub::detail::sub_size_fast_div_mod_impl(ext)...}; +} + +// precompute modulo/division for each mdspan extent +template +[[nodiscard]] _CCCL_HOST_DEVICE_API auto +extents_fast_div_mod(const ::cuda::std::extents& ext, ::cuda::std::index_sequence = {}) +{ + using fast_mod_div_t = fast_div_mod; + using array_t = ::cuda::std::array; + return array_t{fast_mod_div_t(ext.extent(Positions))...}; +} + +// GCC <= 9 constexpr workaround: Extent must be passed as type only, even const Extent& doesn't work +template +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr bool are_extents_in_range_static(int start, int end) +{ + for (auto i = start; i < end; i++) + { + if (Extents::static_extent(i) == ::cuda::std::dynamic_extent) + { + return false; + } + } + return true; +} + +template +[[nodiscard]] _CCCL_HOST_DEVICE_API bool +have_same_strides(const MappingTypeLhs& mapping_lhs, const MappingTypeRhs& mapping_rhs) +{ + auto extents_lhs = mapping_lhs.extents(); + auto extents_rhs = mapping_rhs.extents(); + _CCCL_ASSERT(extents_lhs.rank() == extents_rhs.rank(), "extents must have the same rank"); + for (size_t i = 0; i < extents_lhs.rank(); i++) + { + if (mapping_lhs.stride(i) != mapping_rhs.stride(i)) + { + return false; + } + } + return true; +} +} // namespace detail +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/rfa.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/rfa.cuh new file mode 100644 index 00000000..4fab3864 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/rfa.cuh @@ -0,0 +1,691 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::rfa +{ +// jump table for indexing into data +inline constexpr int cub_rfa_max_jump = 5; +static_assert(cub_rfa_max_jump <= 5, "cub_rfa_max_jump must be less than or equal to 5"); + +template +static _CCCL_DEVICE FType* get_shared_bin_array() +{ + static __shared__ FType bin_computed_array[Len]; + return bin_computed_array; +} + +//! Class to hold a reproducible summation of the numbers passed to it +//! +//! @param FType Floating-point data type; either `float` or `double +//! @param Fold Number of collectors in the binned number (K-fold), used for reproducible summation. Defaults to 3. +template >* = nullptr> +class alignas(2 * sizeof(FType)) ReproducibleFloatingAccumulator +{ +public: + using ftype = FType; + +private: + ::cuda::std::array data{}; + + /// Floating-point precision bin width + static constexpr int bin_width = ::cuda::std::is_same_v ? 40 : 13; + static constexpr int min_exp = ::cuda::std::numeric_limits::min_exponent; + static constexpr int max_exp = ::cuda::std::numeric_limits::max_exponent; + static constexpr int mant_dig = ::cuda::std::numeric_limits::digits; + +public: + /// Binned floating-point maximum index + static constexpr int max_index = ((max_exp - min_exp + mant_dig - 1) / bin_width) - 1; + + // The maximum floating-point fold supported by the library + static constexpr auto max_fold = max_index + 1; + + _CCCL_DEVICE static ftype initialize_bin(int index) noexcept + { + if (index == 0) + { + if constexpr (::cuda::std::is_same_v) + { + return ::cuda::std::ldexp(0.75, max_exp); + } + else + { + return 2.0 * ::cuda::std::ldexp(0.75, max_exp - 1); + } + } + + if (index > 0 && index <= max_index) + { + return ::cuda::std::ldexp(0.75, max_exp + mant_dig - bin_width + 1 - index * bin_width); + } + else + { + return ::cuda::std::ldexp(0.75, max_exp + mant_dig - bin_width + 1 - max_index * bin_width); + } + } + +private: + /// Binned floating-point compression factor + /// This factor is used to scale down inputs before deposition into the bin of + /// highest index + static constexpr auto compression = 1.0 / (1 << (mant_dig - bin_width + 1)); + /// Binned double precision expansion factor + /// This factor is used to scale up inputs after deposition into the bin of + /// highest index + static constexpr auto expansion = 1.0 * (1 << (mant_dig - bin_width + 1)); + static constexpr auto exp_bias = max_exp - 2; + + /// Return a binned floating-point bin + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static ftype binned_bins(int index) + { + ftype* bins = get_shared_bin_array(); + return bins[index]; + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint32_t& get_bit_representation(float& x) noexcept + { + return *reinterpret_cast(&x); + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint64_t& get_bit_representation(double& x) noexcept + { + return *reinterpret_cast(&x); + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint32_t get_bit_representation(const float& x) noexcept + { + return ::cuda::std::bit_cast(x); + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static uint64_t get_bit_representation(const double& x) noexcept + { + return ::cuda::std::bit_cast(x); + } + + /// Return primary vector value const ref + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE const ftype& primary(int i) const noexcept + { + if constexpr (Fold <= cub_rfa_max_jump) + { + switch (i) + { + case 0: + if constexpr (Fold >= 1) + { + return data[0]; + } + [[fallthrough]]; + case 1: + if constexpr (Fold >= 2) + { + return data[1]; + } + [[fallthrough]]; + case 2: + if constexpr (Fold >= 3) + { + return data[2]; + } + [[fallthrough]]; + case 3: + if constexpr (Fold >= 4) + { + return data[3]; + } + [[fallthrough]]; + case 4: + if constexpr (Fold >= 5) + { + return data[4]; + } + [[fallthrough]]; + default: + return data[Fold - 1]; + } + } + else + { + return data[i]; + } + } + + /// Return carry vector value const ref + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE const ftype& carry(int i) const noexcept + { + if (Fold <= cub_rfa_max_jump) + { + switch (i) + { + case 0: + if (Fold >= 1) + { + return data[Fold + 0]; + } + [[fallthrough]]; + case 1: + if (Fold >= 2) + { + return data[Fold + 1]; + } + [[fallthrough]]; + case 2: + if (Fold >= 3) + { + return data[Fold + 2]; + } + [[fallthrough]]; + case 3: + if (Fold >= 4) + { + return data[Fold + 3]; + } + [[fallthrough]]; + case 4: + if (Fold >= 5) + { + return data[Fold + 4]; + } + [[fallthrough]]; + default: + return data[2 * Fold - 1]; + } + } + else + { + return data[Fold + i]; // NOLINT(bugprone-misplaced-widening-cast) + } + } + + /// Return primary vector value ref + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ftype& primary(int i) noexcept + { + const auto& c = *this; + return const_cast(c.primary(i)); + } + + /// Return carry vector value ref + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE ftype& carry(int i) noexcept + { + const auto& c = *this; + return const_cast(c.carry(i)); + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static int exp_val(const ftype x) noexcept + { + const auto bits = get_bit_representation(x); + return (bits >> (mant_dig - 1)) & (2 * max_exp - 1); + } + + /// Get index of float-point precision + /// The index of a non-binned type is the smallest index a binned type would + /// need to have to sum it reproducibly. Higher indices correspond to smaller + /// bins. + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE static int binned_dindex(const ftype x) + { + int exp = exp_val(x); + + if (exp != 0) + { + return ((max_exp + exp_bias) - exp) / bin_width; + } + if (x == 0.0) + { + return max_index; + } + else + { + (void) ::cuda::std::frexpf(x, &exp); + return (::cuda::std::min) ((max_exp - exp) / bin_width, +max_index); + } + } + + /// Get index of manually specified binned double precision + /// The index of a binned type is the bin that it corresponds to. Higher + /// indices correspond to smaller bins. + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE int binned_index() const + { + return ((max_exp + mant_dig - bin_width + 1 + exp_bias) - exp_val(primary(0))) / bin_width; + } + + /// Check if index of manually specified binned floating-point is 0 + /// A quick check to determine if the index is 0 + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE bool is_binned_index_zero() const + { + return exp_val(primary(0)) == max_exp + exp_bias; + } + + //! Update manually specified binned fp with a scalar (X -> Y) + //! + //! This method updates the binned fp to an index suitable for adding numbers + //! with absolute value less than @p max_abs_val + _CCCL_DEVICE void binned_update(const ftype max_abs_val) + { + int X_index = binned_dindex(max_abs_val); + int shift = binned_index() - X_index; + if (shift > 0) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = Fold - 1; i >= 1; i--) + { + if (i < shift) + { + break; + } + primary(i) = primary((i - shift)); + carry(i) = carry((i - shift)); + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < Fold; j++) + { + if (j >= shift) + { + break; + } + primary(j) = binned_bins(j + X_index); + carry(j) = 0.0; + } + } + } + + //! Add scalar @p X to suitably binned manually specified binned fp (Y += X) + //! + //! Performs the operation Y += X on an binned type Y where the index of Y is + //! larger than the index of @p X + _CCCL_DEVICE void binned_deposit(const ftype X) + { + ftype M; + ftype x = X; + + if (is_binned_index_zero()) + { + M = primary(0); + ftype qd = x * compression; + auto& ql = get_bit_representation(qd); + ql |= 1; + qd += M; + primary(0) = qd; + M -= qd; + M *= expansion * 0.5; + x += M; + x += M; + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 1; i < Fold - 1; i++) + { + M = primary(i); + qd = x; + ql |= 1; + qd += M; + primary(i) = qd; + M -= qd; + x += M; + } + qd = x; + ql |= 1; + primary((Fold - 1)) += qd; + } + else + { + ftype qd = x; + auto& ql = get_bit_representation(qd); + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < Fold - 1; i++) + { + M = primary(i); + qd = x; + ql |= 1; + qd += M; + primary(i) = qd; + M -= qd; + x += M; + } + qd = x; + ql |= 1; + primary((Fold - 1)) += qd; + } + } + + //! Renormalize manually specified binned double precision + //! + //! Renormalization keeps the primary vector within the necessary bins by + //! shifting over to the carry vector + _CCCL_DEVICE _CCCL_FORCEINLINE void binned_renorm() + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < Fold; i++) + { + auto tmp_renormd = primary(i); + auto& tmp_renorml = get_bit_representation(tmp_renormd); + + carry(i) += static_cast((tmp_renorml >> (mant_dig - 3)) & 3) - 2; + + tmp_renorml &= ~(1ull << (mant_dig - 3)); + tmp_renorml |= 1ull << (mant_dig - 2); + primary(i) = tmp_renormd; + } + } + + //! Add scalar to manually specified binned fp (Y += X) + //! + //! Performs the operation Y += X on an binned type Y + _CCCL_DEVICE _CCCL_FORCEINLINE void binned_add(const ftype x) + { + binned_update(x); + binned_deposit(x); + binned_renorm(); + } + + //! Add two manually specified binned fp (Y += X) + //! Performs the operation Y += X + //! + //! @param x Another binned fp of the same type + _CCCL_DEVICE void binned_add(const ReproducibleFloatingAccumulator& x) + { + const auto X_index = x.binned_index(); + const auto Y_index = this->binned_index(); + const auto shift = Y_index - X_index; + if (shift > 0) + { + // shift Y upwards and add X to Y + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = Fold - 1; i >= 1; i--) + { + if (i < shift) + { + break; + } + primary(i) = x.primary(i) + (primary((i - shift)) - binned_bins(i - shift + Y_index)); + carry(i) = x.carry(i) + carry((i - shift)); + } + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < Fold; i++) + { + if (i == shift) + { + break; + } + primary(i) = x.primary(i); + carry(i) = x.carry(i); + } + } + else if (shift < 0) + { + // shift X upwards and add X to Y + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < Fold; i++) + { + if (i < -shift) + { + continue; + } + primary(i) += x.primary((i + shift)) - binned_bins(X_index + i + shift); + carry(i) += x.carry((i + shift)); + } + } + else if (shift == 0) + { + // add X to Y + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < Fold; i++) + { + primary(i) += x.primary(i) - binned_bins(i + X_index); + carry(i) += x.carry(i); + } + } + + binned_renorm(); + } + + [[nodiscard]] _CCCL_DEVICE double conv_binned_to_double() const + { + int i = 0; + double Y = 0.0; + const auto X_index = binned_index(); + if (X_index <= (3 * mant_dig) / bin_width) + { + double scale_down = ::cuda::std::ldexpf(0.5f, 1 - (2 * mant_dig - bin_width)); + double scale_up = ::cuda::std::ldexpf(0.5f, 1 - (2 * mant_dig - bin_width)); + int scaled = ::cuda::std::max(::cuda::std::min(Fold, (3 * mant_dig) / bin_width - X_index), 0); + if (X_index == 0) + { + Y += carry(0) * ((binned_bins(0 + X_index) / 6.0) * scale_down * expansion); + Y += carry(1) * ((binned_bins(1 + X_index) / 6.0) * scale_down); + Y += (primary(0) - binned_bins(0 + X_index)) * scale_down * expansion; + i = 2; + } + else + { + Y += carry(0) * ((binned_bins(0 + X_index) / 6.0) * scale_down); + i = 1; + } + for (; i < scaled; i++) + { + Y += carry(i) * ((binned_bins(i + X_index) / 6.0) * scale_down); + Y += (primary((i - 1)) - binned_bins(i - 1 + X_index)) * scale_down; + } + if (i == Fold) + { + Y += (primary((Fold - 1)) - binned_bins(Fold - 1 + X_index)) * scale_down; + return Y * scale_up; + } + if (::cuda::std::isinf(Y * scale_up)) + { + return Y * scale_up; + } + Y *= scale_up; + for (; i < Fold; i++) + { + Y += carry(i) * (binned_bins(i + X_index) / 6.0); + Y += primary((i - 1)) - binned_bins(i - 1 + X_index); + } + Y += primary((Fold - 1)) - binned_bins(Fold - 1 + X_index); + } + else + { + Y += carry(0) * (binned_bins(0 + X_index) / 6.0); + for (i = 1; i < Fold; i++) + { + Y += carry(i) * (binned_bins(i + X_index) / 6.0); + Y += (primary((i - 1)) - binned_bins(i - 1 + X_index)); + } + Y += (primary((Fold - 1)) - binned_bins(Fold - 1 + X_index)); + } + return Y; + } + + [[nodiscard]] _CCCL_DEVICE float conv_binned_to_float() const + { + int i = 0; + double Y = 0.0; + + // Note that the following order of summation is in order of decreasing + // exponent. The following code is specific to SBWIDTH=13, FLT_MANT_DIG=24, and + // the number of carries equal to 1. + const auto X_index = binned_index(); + if (X_index == 0) + { + Y += static_cast(carry(0)) * static_cast(binned_bins(0 + X_index) / 6.0) + * static_cast(expansion); + Y += static_cast(carry(1)) * static_cast(binned_bins(1 + X_index) / 6.0); + Y += static_cast(primary(0) - binned_bins(0 + X_index)) * static_cast(expansion); + i = 2; + } + else + { + Y += static_cast(carry(0)) * static_cast((binned_bins(0 + X_index) / 6.0)); + i = 1; + } + for (; i < Fold; i++) + { + Y += static_cast(carry(i)) * static_cast(binned_bins(i + X_index) / 6.0); + Y += static_cast(primary(i - 1) - binned_bins(i - 1 + X_index)); + } + Y += static_cast(primary(Fold - 1) - binned_bins(Fold - 1 + X_index)); + return static_cast(Y); + } + +public: + ReproducibleFloatingAccumulator() = default; + + /// Set the binned fp to zero + _CCCL_DEVICE void zero() noexcept + { + data = {}; + } + + [[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE constexpr int endurance() const noexcept + { + return 1 << (mant_dig - bin_width - 2); + } + + //! Accumulate an arithmetic @p x into the binned fp. + //! NOTE: Casts @p x to the type of the binned fp + _CCCL_TEMPLATE(typename U) + _CCCL_REQUIRES(::cuda::std::is_arithmetic_v) + _CCCL_DEVICE ReproducibleFloatingAccumulator& operator+=(const U x) + { + binned_add(static_cast(x)); + return *this; + } + + //! Accumulate-subtract an arithmetic @p x into the binned fp. + //! NOTE: Casts @p x to the type of the binned fp + _CCCL_TEMPLATE(typename U) + _CCCL_REQUIRES(::cuda::std::is_arithmetic_v) + _CCCL_DEVICE ReproducibleFloatingAccumulator& operator-=(const U x) + { + binned_add(-static_cast(x)); + return *this; + } + + /// Accumulate a binned fp @p x into the binned fp. + _CCCL_DEVICE ReproducibleFloatingAccumulator& operator+=(const ReproducibleFloatingAccumulator& other) + { + binned_add(other); + return *this; + } + + //! Accumulate-subtract a binned fp @p other into the binned fp. + //! NOTE: Makes a copy and performs arithmetic; slow. + _CCCL_DEVICE ReproducibleFloatingAccumulator& operator-=(const ReproducibleFloatingAccumulator& other) + { + const auto temp = -other; + binned_add(temp); + } + + _CCCL_DEVICE friend bool operator==(const ReproducibleFloatingAccumulator& a, const ReproducibleFloatingAccumulator& b) + { + return a.data == b.data; + } + + _CCCL_DEVICE friend bool operator!=(const ReproducibleFloatingAccumulator& a, const ReproducibleFloatingAccumulator& b) + { + return !(a == b); + } + + //! Sets this binned fp equal to the arithmetic value @p x + //! NOTE: Casts @p x to the type of the binned fp + _CCCL_TEMPLATE(typename U) + _CCCL_REQUIRES(::cuda::std::is_arithmetic_v) + _CCCL_DEVICE ReproducibleFloatingAccumulator& operator=(const U x) + { + zero(); + binned_add(static_cast(x)); + return *this; + } + + //! Returns the negative of this binned fp + //! NOTE: Makes a copy and performs arithmetic; slow. + [[nodiscard]] _CCCL_DEVICE ReproducibleFloatingAccumulator operator-() const + { + ReproducibleFloatingAccumulator temp = *this; + if (primary(0) != 0.0) + { + _CCCL_PRAGMA_UNROLL_FULL() + for (int i = 0; i < Fold; i++) + { + temp.primary(i) = binned_bins(i + binned_index()) - (primary(i) - binned_bins(i + binned_index())); + temp.carry(i) = -carry(i); + } + } + return temp; + } + + /// Convert this binned fp into its native floating-point representation + [[nodiscard]] _CCCL_DEVICE ftype conv_to_fp() const + { + if (::cuda::std::is_same_v) + { + return conv_binned_to_float(); + } + else + { + return conv_binned_to_double(); + } + } + + /// Add @p x to the binned fp + _CCCL_DEVICE void add(const ftype x) + { + binned_add(x); + } + + ////////////////////////////////////// + // MANUAL OPERATIONS; USE WISELY + ////////////////////////////////////// + + //! Rebins for repeated accumulation of scalars with magnitude <= @p mav + //! + //! Once rebinned, `endurance` values <= @p mav can be added to the accumulator + //! with `unsafe_add` after which `renorm()` must be called. See the source of + //!`add()` for an example + _CCCL_DEVICE void set_max_val(const ftype mav) + { + binned_update(mav); + } + + //! Add @p x to the binned fp + //! + //! This is intended to be used after a call to `set_max_abs_val()` + _CCCL_DEVICE void unsafe_add(const ftype x) + { + binned_deposit(x); + } + + //! Renormalizes the binned fp + //! + //! This is intended to be used after a call to `set_max_abs_val()` and one or + //! more calls to `unsafe_add()` + _CCCL_DEVICE void renorm() + { + binned_renorm(); + } +}; +} // namespace detail::rfa + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/segmented_params.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/segmented_params.cuh new file mode 100644 index 00000000..0ece8092 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/segmented_params.cuh @@ -0,0 +1,145 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::params +{ +// ===================================================================== +// get_param — unified segment parameter access +// ===================================================================== + +//! @brief Returns the value of an argument for a given segment index. +//! +//! @param[in] __arg Argument or argument wrapper to read. +//! @param[in] __index Segment index to read for sequence arguments. +//! @return The single argument value, or the sequence element at the given index. +_CCCL_TEMPLATE(class _Tp, class _SegmentIndexT) +_CCCL_REQUIRES((!::cuda::args::__is_wrapper_v<::cuda::std::remove_cvref_t<_Tp>>) ) +[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto get_param(_Tp&& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept +{ + if constexpr (::cuda::args::__traits<::cuda::std::remove_cvref_t<_Tp>>::is_single_value) + { + return __arg; + } + else + { + return __arg[__index]; + } +} + +template +[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto +get_param(const ::cuda::args::constant<_Value, _Tp>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept +{ + return ::cuda::args::__unwrap(__arg); +} + +template +[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto +get_param(const ::cuda::args::immediate<_Arg, _StaticBounds>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept +{ + return ::cuda::args::__unwrap(__arg); +} + +template +[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto +get_param(const ::cuda::args::deferred<_Arg, _StaticBounds>& __arg, [[maybe_unused]] _SegmentIndexT __index) noexcept +{ + return ::cuda::args::__unwrap(__arg); +} + +template +[[nodiscard]] _CCCL_HOST_DEVICE constexpr auto +get_param(const ::cuda::args::deferred_sequence<_Arg, _StaticBounds>& __arg, _SegmentIndexT __index) noexcept +{ + return ::cuda::args::__unwrap(__arg)[__index]; +} + +// ===================================================================== +// Discrete parameter support +// ===================================================================== + +//! @brief Specifies a list of supported options for a parameter. +template +struct supported_options +{ + static constexpr ::cuda::std::size_t count = sizeof...(Options); +}; + +//! @brief Static discrete parameter — a single compile-time value that is also its only supported option. +//! +//! Holds no runtime value, so it cannot be put into a state that disagrees with its supported option, and +//! @c dispatch_impl therefore always matches it. This is the safe representation for a compile-time-fixed discrete +//! parameter (e.g. a statically known top-k selection direction): modeling such a parameter with a runtime value +//! instead would risk that value silently disagreeing with the supported option (a no-op dispatch unless +//! @c CCCL_ENABLE_ASSERTIONS is set). +template +struct static_discrete_param +{ + using value_type = T; + using supported_options_t = supported_options; + + template + [[nodiscard]] _CCCL_HOST_DEVICE constexpr T get_param(SegmentIndexT) const noexcept + { + return Value; + } +}; + +// ===================================================================== +// Discrete dispatch +// ===================================================================== + +//! @brief Translates a runtime parameter value into a compile-time constant by matching +//! against a list of supported options. +//! +//! @param[in] val Runtime value to match. +//! @param[in] __supported_options Supported values for the parameter. +//! @param[in] f Functor invoked with the matched compile-time constant. +//! @return `true` if the value matches one of the supported options. +template +[[nodiscard]] _CCCL_HOST_DEVICE bool +dispatch_impl(T val, [[maybe_unused]] supported_options __supported_options, Functor&& f) +{ + const bool match_found = ((val == Opts ? (f(::cuda::std::integral_constant{}), true) : false) || ...); + _CCCL_ASSERT(match_found, "The given runtime parameter value is not in the supported list"); + return match_found; +} + +//! @brief Dispatcher that resolves a discrete parameter to a compile-time constant +//! and invokes a functor with the matched option. +//! +//! @param[in] param Discrete parameter to resolve. +//! @param[in] segment_id Segment index to read from `param`. +//! @param[in] f Functor invoked with the matched compile-time constant. +//! @return `true` if the parameter value matches one of its supported options. +template +[[nodiscard]] _CCCL_HOST_DEVICE bool dispatch_discrete(ParamT param, SegmentIndexT segment_id, Functor&& f) +{ + using supported_list = typename ParamT::supported_options_t; + auto param_value = param.get_param(segment_id); + return CUB_NS_QUALIFIER::detail::params::dispatch_impl( + param_value, supported_list{}, ::cuda::std::forward(f)); +} +} // namespace detail::params + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/strong_load.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/strong_load.cuh new file mode 100644 index 00000000..b42c2e53 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/strong_load.cuh @@ -0,0 +1,167 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file Utilities for strong memory operations. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +CUB_NAMESPACE_BEGIN + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + +namespace detail +{ +static _CCCL_DEVICE _CCCL_FORCEINLINE uint4 load_relaxed(uint4 const* ptr) +{ + uint4 retval; + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("ld.relaxed.gpu.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(retval.x), + "=r"(retval.y), + "=r"(retval.z), + "=r"(retval.w) : "l"(ptr) : "memory");), + (asm volatile("ld.cg.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(retval.x), + "=r"(retval.y), + "=r"(retval.z), + "=r"(retval.w) : "l"(ptr) : "memory");)); + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE ulonglong2 load_relaxed(ulonglong2 const* ptr) +{ + ulonglong2 retval; + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("ld.relaxed.gpu.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");), + (asm volatile("ld.cg.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");)); + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE ushort4 load_relaxed(ushort4 const* ptr) +{ + ushort4 retval; + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("ld.relaxed.gpu.v4.u16 {%0, %1, %2, %3}, [%4];" : "=h"(retval.x), + "=h"(retval.y), + "=h"(retval.z), + "=h"(retval.w) : "l"(ptr) : "memory");), + (asm volatile("ld.cg.v4.u16 {%0, %1, %2, %3}, [%4];" : "=h"(retval.x), + "=h"(retval.y), + "=h"(retval.z), + "=h"(retval.w) : "l"(ptr) : "memory");)); + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE uint2 load_relaxed(uint2 const* ptr) +{ + uint2 retval; + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("ld.relaxed.gpu.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");), + (asm volatile("ld.cg.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");)); + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned long long load_relaxed(unsigned long long const* ptr) +{ + unsigned long long retval; + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, + (asm volatile("ld.relaxed.gpu.u64 %0, [%1];" : "=l"(retval) : "l"(ptr) : "memory");), + (asm volatile("ld.cg.u64 %0, [%1];" : "=l"(retval) : "l"(ptr) : "memory");)); + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int load_relaxed(unsigned int const* ptr) +{ + unsigned int retval; + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, + (asm volatile("ld.relaxed.gpu.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");), + (asm volatile("ld.cg.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");)); + + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned short load_relaxed(unsigned short const* ptr) +{ + unsigned short retval; + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, + (asm volatile("ld.relaxed.gpu.u16 %0, [%1];" : "=h"(retval) : "l"(ptr) : "memory");), + (asm volatile("ld.cg.u16 %0, [%1];" : "=h"(retval) : "l"(ptr) : "memory");)); + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned char load_relaxed(unsigned char const* ptr) +{ + unsigned short retval; + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("{" + " .reg .u8 datum;" + " ld.relaxed.gpu.u8 datum, [%1];" + " cvt.u16.u8 %0, datum;" + "}" : "=h"(retval) : "l"(ptr) : "memory");), + (asm volatile("{" + " .reg .u8 datum;" + " ld.cg.u8 datum, [%1];" + " cvt.u16.u8 %0, datum;" + "}" : "=h"(retval) : "l"(ptr) : "memory");)); + return (unsigned char) retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE ulonglong2 load_acquire(ulonglong2 const* ptr) +{ + ulonglong2 retval; + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("ld.acquire.gpu.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory");), + ({ + asm volatile("ld.cg.v2.u64 {%0, %1}, [%2];" : "=l"(retval.x), "=l"(retval.y) : "l"(ptr) : "memory"); + __threadfence(); + })); + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE uint2 load_acquire(uint2 const* ptr) +{ + uint2 retval; + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("ld.acquire.gpu.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory");), + ({ + asm volatile("ld.cg.v2.u32 {%0, %1}, [%2];" : "=r"(retval.x), "=r"(retval.y) : "l"(ptr) : "memory"); + __threadfence(); + })); + return retval; +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE unsigned int load_acquire(unsigned int const* ptr) +{ + unsigned int retval; + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, + (asm volatile("ld.acquire.gpu.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory");), + (asm volatile("ld.cg.u32 %0, [%1];" : "=r"(retval) : "l"(ptr) : "memory"); __threadfence();)); + + return retval; +} +} // namespace detail + +#endif // _CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/strong_store.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/strong_store.cuh new file mode 100644 index 00000000..086b4802 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/strong_store.cuh @@ -0,0 +1,223 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file Utilities for strong memory operations. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +CUB_NAMESPACE_BEGIN + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + +namespace detail +{ +static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(uint4* ptr, uint4 val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("st.relaxed.gpu.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr), + "r"(val.x), + "r"(val.y), + "r"(val.z), + "r"(val.w) : "memory");), + (asm volatile( + "st.cg.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w) : "memory");)); +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(ulonglong2* ptr, ulonglong2 val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("st.relaxed.gpu.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");), + (asm volatile("st.cg.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");)); +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(ushort4* ptr, ushort4 val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("st.relaxed.gpu.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr), + "h"(val.x), + "h"(val.y), + "h"(val.z), + "h"(val.w) : "memory");), + (asm volatile( + "st.cg.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr), "h"(val.x), "h"(val.y), "h"(val.z), "h"(val.w) : "memory");)); +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(uint2* ptr, uint2 val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("st.relaxed.gpu.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");), + (asm volatile("st.cg.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");)); +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned long long* ptr, unsigned long long val) +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, + (asm volatile("st.relaxed.gpu.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");), + (asm volatile("st.cg.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");)); +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned int* ptr, unsigned int val) +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, + (asm volatile("st.relaxed.gpu.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");), + (asm volatile("st.cg.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");)); +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned short* ptr, unsigned short val) +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_70, + (asm volatile("st.relaxed.gpu.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");), + (asm volatile("st.cg.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");)); +} + +static _CCCL_DEVICE _CCCL_FORCEINLINE void store_relaxed(unsigned char* ptr, unsigned char val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("{" + " .reg .u8 datum;" + " cvt.u8.u16 datum, %1;" + " st.relaxed.gpu.u8 [%0], datum;" + "}" : : "l"(ptr), + "h"((unsigned short) val) : "memory");), + (asm volatile("{" + " .reg .u8 datum;" + " cvt.u8.u16 datum, %1;" + " st.cg.u8 [%0], datum;" + "}" : : "l"(ptr), + "h"((unsigned short) val) : "memory");)); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(uint4* ptr, uint4 val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("st.release.gpu.v4.u32 [%0], {%1, %2, %3, %4};" : : "l"(ptr), + "r"(val.x), + "r"(val.y), + "r"(val.z), + "r"(val.w) : "memory");), + ({ + __threadfence(); + asm volatile("st.cg.v4.u32 [%0], {%1, %2, %3, %4};" + : + : "l"(ptr), "r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w) + : "memory"); + })); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(ulonglong2* ptr, ulonglong2 val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("st.release.gpu.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory");), + ({ + __threadfence(); + asm volatile("st.cg.v2.u64 [%0], {%1, %2};" : : "l"(ptr), "l"(val.x), "l"(val.y) : "memory"); + })); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(ushort4* ptr, ushort4 val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("st.release.gpu.v4.u16 [%0], {%1, %2, %3, %4};" : : "l"(ptr), + "h"(val.x), + "h"(val.y), + "h"(val.z), + "h"(val.w) : "memory");), + ({ + __threadfence(); + asm volatile("st.cg.v4.u16 [%0], {%1, %2, %3, %4};" + : + : "l"(ptr), "h"(val.x), "h"(val.y), "h"(val.z), "h"(val.w) + : "memory"); + })); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(uint2* ptr, uint2 val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("st.release.gpu.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory");), + ({ + __threadfence(); + asm volatile("st.cg.v2.u32 [%0], {%1, %2};" : : "l"(ptr), "r"(val.x), "r"(val.y) : "memory"); + })); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned long long* ptr, unsigned long long val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory");), ({ + __threadfence(); + asm volatile("st.cg.u64 [%0], %1;" : : "l"(ptr), "l"(val) : "memory"); + })); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned int* ptr, unsigned int val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory");), ({ + __threadfence(); + asm volatile("st.cg.u32 [%0], %1;" : : "l"(ptr), "r"(val) : "memory"); + })); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned short* ptr, unsigned short val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, (asm volatile("st.release.gpu.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory");), ({ + __threadfence(); + asm volatile("st.cg.u16 [%0], %1;" : : "l"(ptr), "h"(val) : "memory"); + })); +} + +_CCCL_DEVICE _CCCL_FORCEINLINE void store_release(unsigned char* ptr, unsigned char val) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_70, + (asm volatile("{" + " .reg .u8 datum;" + " cvt.u8.u16 datum, %1;" + " st.release.gpu.u8 [%0], datum;" + "}" : : "l"(ptr), + "h"((unsigned short) val) : "memory");), + ({ + __threadfence(); + asm volatile( + "{" + " .reg .u8 datum;" + " cvt.u8.u16 datum, %1;" + " st.cg.u8 [%0], datum;" + "}" + : + : "l"(ptr), "h"((unsigned short) val) + : "memory"); + })); +} +} // namespace detail + +#endif // _CCCL_DOXYGEN_INVOKED + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/temporary_storage.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/temporary_storage.cuh new file mode 100644 index 00000000..fe9fca88 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/temporary_storage.cuh @@ -0,0 +1,380 @@ +/* + * Copyright 2021 NVIDIA Corporation + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::temporary_storage +{ +class slot; + +template +class alias; + +template +class layout; + +/** + * @brief Temporary storage slot that can be considered a C++ union with an + * arbitrary fields count. + * + * @warning slot lifetime is defined by the lifetime of the associated layout. + * It's impossible to request new array if layout is already mapped. + * + * @par A Simple Example + * @code + * auto slot = temporary_storage.get_slot(0); + * + * // Add fields into the slot + * // Create an int alias with 0 elements: + * auto int_array = slot->create_alias(); + * // Create a double alias with 2 elements: + * auto double_array = slot->create_alias(2); + * // Create a char alias with 0 elements: + * auto empty_array = slot->create_alias(); + * // Slot size is defined by double_array size (2 * sizeof(double)) + * + * if (condition) + * { + * int_array.grow(42); + * // Now slot size is defined by int_array size (42 * sizeof(int)) + * } + * + * // Temporary storage mapping + * // ... + + * int *d_int_array = int_array.get(); + * double *d_double_array = double_array.get(); + * char *d_empty_array = empty_array.get(); // Guaranteed to return nullptr + * @endcode + */ +class slot +{ + size_t m_size{}; + void* m_pointer{}; + +public: + slot() = default; + + /** + * @brief Returns an array of type @p T and length @p elements + */ + template + _CCCL_HOST_DEVICE alias create_alias(size_t elements = 0); + +private: + _CCCL_HOST_DEVICE void set_bytes_required(size_t new_size) + { + m_size = (::cuda::std::max) (m_size, new_size); + } + + _CCCL_HOST_DEVICE size_t get_bytes_required() const + { + return m_size; + } + + _CCCL_HOST_DEVICE void set_storage(void* ptr) + { + m_pointer = ptr; + } + _CCCL_HOST_DEVICE void* get_storage() const + { + return m_pointer; + } + + template + friend class alias; + + template + friend class layout; +}; + +/** + * @brief Named memory region of a temporary storage slot + * + * @par Overview + * This class provides a typed wrapper of a temporary slot memory region. + * It can be considered as a field in the C++ union. It's only possible to + * increase the array size. + * + * @warning alias lifetime is defined by the lifetime of the associated slot + * It's impossible to grow the array if the layout is already mapped. + */ +template +class alias +{ + slot& m_slot; + size_t m_elements{}; + + _CCCL_HOST_DEVICE explicit alias(slot& slot, size_t elements = 0) + : m_slot(slot) + , m_elements(elements) + { + this->update_slot(); + } + + _CCCL_HOST_DEVICE void update_slot() + { + m_slot.set_bytes_required(m_elements * sizeof(T)); + } + +public: + alias() = delete; + + /** + * @brief Increases the number of elements + * + * @warning + * This method should be called before temporary storage mapping stage. + * + * @param[in] new_elements Increases the memory region occupied in the + * temporary slot to fit up to @p new_elements items + * of type @p T. + */ + _CCCL_HOST_DEVICE void grow(size_t new_elements) + { + m_elements = new_elements; + this->update_slot(); + } + + /** + * @brief Returns pointer to array + * + * If the @p elements number is equal to zero, or storage layout isn't mapped, + * @p nullptr is returned. + */ + _CCCL_HOST_DEVICE T* get() const + { + if (m_elements == 0) + { + return nullptr; + } + + return reinterpret_cast(m_slot.get_storage()); + } + + friend class slot; +}; + +template +_CCCL_HOST_DEVICE alias slot::create_alias(size_t elements) +{ + return alias(*this, elements); +} + +/** + * @brief Temporary storage layout represents a structure with + * @p SlotsCount union-like fields + * + * The layout can be mapped to a temporary buffer only once. + * + * @par A Simple Example + * @code + * cub::detail::temporary_storage::layout<3> temporary_storage; + * + * auto slot_1 = temporary_storage.get_slot(0); + * auto slot_2 = temporary_storage.get_slot(1); + * + * // Add fields into the first slot + * auto int_array = slot_1->create_alias(1); + * auto double_array = slot_1->create_alias(2); + * + * // Add fields into the second slot + * auto char_array = slot_2->create_alias(); + * + * // The equivalent C++ structure could look like + * // struct StorageLayout + * // { + * // union { + * // } slot_0; + * // std::byte padding_0[256 - sizeof (slot_0)]; + * // + * // union { + * // int alias_0[1]; + * // double alias_1[2]; + * // } slot_1; + * // std::byte padding_1[256 - sizeof (slot_1)]; + * // + * // union { + * // char alias_0[0]; + * // } slot_2; + * // std::byte padding_2[256 - sizeof (slot_2)]; + * // }; + * + * // The third slot is empty + * + * // Temporary storage mapping + * if (d_temp_storage == nullptr) + * { + * temp_storage_bytes = temporary_storage.get_size(); + * return; + * } + * else + * { + * temporary_storage.map_to_buffer(d_temp_storage, temp_storage_bytes); + * } + * + * // Use pointers + * int *d_int_array = int_array.get(); + * double *d_double_array = double_array.get(); + * char *d_char_array = char_array.get(); + * @endcode + */ +template +class layout +{ + slot m_slots[SlotsCount]; + size_t m_sizes[SlotsCount]; + void* m_pointers[SlotsCount]; + bool m_layout_was_mapped{}; + +public: + layout() = default; + + _CCCL_HOST_DEVICE slot* get_slot(int slot_id) + { + if (slot_id < SlotsCount) + { + return &m_slots[slot_id]; + } + + return nullptr; + } + + /** + * @brief Returns required temporary storage size in bytes + */ + _CCCL_HOST_DEVICE size_t get_size() + { + this->prepare_interface(); + + // alias_temporaries can return error only in mapping stage, so it's safe to ignore it here. + size_t temp_storage_bytes{}; + [[maybe_unused]] const auto error = detail::alias_temporaries(nullptr, temp_storage_bytes, m_pointers, m_sizes); + _CCCL_ASSERT(error == cudaSuccess, ""); + + if (temp_storage_bytes == 0) + { + // The current CUB convention implies that there are two stages for each + // device-scope function call. The first one returns the required storage + // size. The second stage consumes temporary storage to perform some work. + // The only way to distinguish between the two stages is by checking the + // value of the temporary storage pointer. If zero bytes are requested, + // `cudaMalloc` will return `nullptr`. This fact makes it impossible to + // distinguish between the two stages, so we request some fixed amount of + // bytes (even if we don't need it) to have a non-null temporary storage + // pointer. + return 1; + } + + return temp_storage_bytes; + } + + /** + * @brief Maps the layout to the temporary storage buffer. + */ + _CCCL_HOST_DEVICE cudaError_t map_to_buffer(void* d_temp_storage, size_t temp_storage_bytes) + { + if (m_layout_was_mapped) + { + return cudaErrorAlreadyMapped; + } + + this->prepare_interface(); + + if (cudaError_t error = detail::alias_temporaries(d_temp_storage, temp_storage_bytes, m_pointers, m_sizes)) + { + return error; + } + + for (size_t slot_id = 0; slot_id < SlotsCount; slot_id++) + { + m_slots[slot_id].set_storage(m_pointers[slot_id]); + } + + m_layout_was_mapped = true; + return cudaSuccess; + } + +private: + _CCCL_HOST_DEVICE void prepare_interface() + { + if (m_layout_was_mapped) + { + return; + } + + for (size_t slot_id = 0; slot_id < SlotsCount; slot_id++) + { + const size_t slot_size = m_slots[slot_id].get_bytes_required(); + + m_sizes[slot_id] = slot_size; + m_pointers[slot_id] = nullptr; + } + } +}; + +template +CUB_RUNTIME_FUNCTION cudaError_t +allocate(::cuda::stream_ref stream, void*& d_temp_storage, size_t temp_storage_bytes, MRT& mr) +{ + _CCCL_TRY + { + d_temp_storage = mr.allocate(stream, temp_storage_bytes, alignof(::cuda::std::max_align_t)); + return cudaSuccess; + } + _CCCL_CATCH_ALL + { + return cudaErrorMemoryAllocation; + } +} + +template +CUB_RUNTIME_FUNCTION cudaError_t +deallocate(::cuda::stream_ref stream, void* d_temp_storage, size_t temp_storage_bytes, MRT& mr) +{ + _CCCL_TRY + { + mr.deallocate(stream, d_temp_storage, temp_storage_bytes, alignof(::cuda::std::max_align_t)); + return cudaSuccess; + } + _CCCL_CATCH_ALL + { + return cudaErrorMemoryAllocation; + } +} +} // namespace detail::temporary_storage + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/type_traits.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/type_traits.cuh new file mode 100644 index 00000000..705f71ca --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/type_traits.cuh @@ -0,0 +1,185 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +/** + * \file + * Wrappers and extensions around utilities. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include // IWYU pragma: keep +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +template +inline constexpr bool is_one_of_v = (::cuda::std::is_same_v || ...); + +template +struct has_binary_call_operator : ::cuda::std::false_type +{}; + +template +struct has_binary_call_operator< + T, + V, + ::cuda::std::void_t()(::cuda::std::declval(), ::cuda::std::declval()))>> + : ::cuda::std::true_type +{}; + +/*********************************************************************************************************************** + * Array-like type traits + **********************************************************************************************************************/ + +template +inline constexpr bool is_fixed_size_random_access_range_v = false; + +template +inline constexpr bool is_fixed_size_random_access_range_v = true; + +template +inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::array> = true; + +template +inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::span> = N != ::cuda::std::dynamic_extent; + +template +inline constexpr bool is_fixed_size_random_access_range_v<::cuda::std::mdspan> = + E::rank() == 1 && E::rank_dynamic() == 0; + +/*********************************************************************************************************************** + * static_size: a type trait that returns the number of elements in an Array-like type + **********************************************************************************************************************/ + +template +inline constexpr int static_size_v = ::cuda::std::enable_if_t<::cuda::std::__always_false_v>{}; + +template +inline constexpr int static_size_v = N; + +template +inline constexpr int static_size_v<::cuda::std::array> = N; + +template +inline constexpr int static_size_v<::cuda::std::span> = + ::cuda::std::enable_if_t{N}; + +template +inline constexpr int static_size_v<::cuda::std::mdspan> = + ::cuda::std::enable_if_t{E::static_extent(0)}; + +template +using implicit_prom_t = decltype(+T{}); + +/*********************************************************************************************************************** + * Extended floating point traits + **********************************************************************************************************************/ +// half + +template +inline constexpr bool is_half_impl_v = false; + +template +inline constexpr bool is_half2_impl_v = false; + +#if _CCCL_HAS_NVFP16() + +template <> +inline constexpr bool is_half_impl_v<__half> = true; + +template <> +inline constexpr bool is_half2_impl_v<__half2> = true; + +#endif // _CCCL_HAS_NVFP16 + +template +inline constexpr bool is_half_v = is_half_impl_v<::cuda::std::remove_cv_t>; + +template +inline constexpr bool is_half2_v = is_half2_impl_v<::cuda::std::remove_cv_t>; + +template +inline constexpr bool is_any_half_v = is_half_impl_v || is_half2_impl_v; + +//---------------------------------------------------------------------------------------------------------------------- +// bfloat16 + +template +inline constexpr bool is_bfloat16_impl_v = false; + +template +inline constexpr bool is_bfloat162_impl_v = false; + +#if _CCCL_HAS_NVBF16() + +template <> +inline constexpr bool is_bfloat16_impl_v<__nv_bfloat16> = true; + +template <> +inline constexpr bool is_bfloat162_impl_v<__nv_bfloat162> = true; + +#endif // _CCCL_HAS_NVBF16 + +template +inline constexpr bool is_bfloat16_v = is_bfloat16_impl_v<::cuda::std::remove_cv_t>; + +template +inline constexpr bool is_bfloat162_v = is_bfloat162_impl_v<::cuda::std::remove_cv_t>; + +template +inline constexpr bool is_any_bfloat16_v = is_bfloat16_v || is_bfloat162_v; + +//---------------------------------------------------------------------------------------------------------------------- +// short2/ushort2 + +template +inline constexpr bool is_any_short2_impl_v = false; + +template <> +inline constexpr bool is_any_short2_impl_v = true; + +template <> +inline constexpr bool is_any_short2_impl_v = true; + +template +inline constexpr bool is_any_short2_v = is_any_short2_impl_v<::cuda::std::remove_cv_t>; + +//---------------------------------------------------------------------------------------------------------------------- + +// - promote small integer types to their corresponding 32-bit promotion type +// - address the incompatibility between linux/windows for int/long +template +using signed_promotion_t = ::cuda::std::conditional_t< + ::cuda::std::__cccl_is_signed_integer_v && sizeof(T) <= sizeof(int), + int, + ::cuda::std::conditional_t<::cuda::std::__cccl_is_unsigned_integer_v && sizeof(T) <= sizeof(uint32_t), uint32_t, T>>; +} // namespace detail +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/uninitialized_copy.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/uninitialized_copy.cuh new file mode 100644 index 00000000..253aea8d --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/uninitialized_copy.cuh @@ -0,0 +1,48 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +#if _CCCL_CUDA_COMPILER(NVHPC) +template +_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val) +{ + // NVBug 3384810 + new (ptr) T(::cuda::std::forward(val)); +} +#else // ^^^ _CCCL_CUDA_COMPILER(NVHPC) ^^^ / vvv !_CCCL_CUDA_COMPILER(NVHPC) vvv +template , int> = 0> +_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val) +{ + // gevtushenko: placement new should work here as well, but the code generated for copy assignment is sometimes better + *ptr = ::cuda::std::forward(val); +} + +template , int> = 0> +_CCCL_HOST_DEVICE void uninitialized_copy_single(T* ptr, U&& val) +{ + new (ptr) T(::cuda::std::forward(val)); +} +#endif // !_CCCL_CUDA_COMPILER(NVHPC) +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/unsafe_bitcast.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/unsafe_bitcast.cuh new file mode 100644 index 00000000..cf20deb1 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/unsafe_bitcast.cuh @@ -0,0 +1,34 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +CUB_NAMESPACE_BEGIN +namespace detail +{ +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + +// NOTE: bit_cast cannot be always used because __half, __nv_bfloat16, etc. are not trivially copyable +template +[[nodiscard]] _CCCL_DEVICE _CCCL_FORCEINLINE Output unsafe_bitcast(const Input& input) +{ + Output output; + static_assert(sizeof(input) == sizeof(output), "wrong size"); + // NOLINTNEXTLINE(bugprone-undefined-memory-manipulation) + ::memcpy(&output, &input, sizeof(input)); + return output; +} + +#endif // !_CCCL_DOXYGEN_INVOKED +} // namespace detail +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/allocators/smem_allocator.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/allocators/smem_allocator.cuh new file mode 100644 index 00000000..1fe10bad --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/allocators/smem_allocator.cuh @@ -0,0 +1,83 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +struct SmemAllocator +{ + ::cuda::std::uint32_t mPtrSmem32 = 0; + int mAllocatedSize = 0; + + _CCCL_HOST_DEVICE_API constexpr SmemAllocator() noexcept + { + // we only need the real pointer at runtime in device code + if (!::cuda::std::is_constant_evaluated()) + { + NV_IF_TARGET(NV_IS_DEVICE, mPtrSmem32 = dynamic_smem_base();) + } + } + + [[nodiscard]] _CCCL_DEVICE_API static ::cuda::std::uint32_t dynamic_smem_base() noexcept + { + extern __shared__ char warpSpeedDynamicSmemBase[]; + return __cvta_generic_to_shared(warpSpeedDynamicSmemBase); + } + + // SmemAllocator is a non-copyable, non-movable type. It must be passed by + // (mutable) reference to be useful. + SmemAllocator(const SmemAllocator&) = delete; // Delete copy constructor + SmemAllocator(SmemAllocator&&) = delete; // Delete move constructor + SmemAllocator& operator=(const SmemAllocator&) = delete; // Delete copy assignment + SmemAllocator& operator=(SmemAllocator&&) = delete; // Delete move assignment + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr void* alloc(::cuda::std::uint32_t size, ::cuda::std::uint32_t align = 0) + { + // Align mPtrSmem32 to requested alignment (round-up) + ::cuda::std::uint32_t ptrAllocation32 = (mPtrSmem32 + (align - 1)) & ~(align - 1); + + // Move base pointer and update allocated size + mAllocatedSize += static_cast(size + ptrAllocation32 - mPtrSmem32); + mPtrSmem32 = ptrAllocation32 + size; + + // we only need the pointer at runtime in device code + if (!::cuda::std::is_constant_evaluated()) + { + NV_IF_TARGET( + NV_IS_DEVICE, + ( + // Convert allocated smem address to generic pointer + void* mPtrAllocation = __cvta_shared_to_generic(ptrAllocation32); + // Ensure alignment calculation does not move down into rest of kernel code. + return optimizeSmemPtr(mPtrAllocation);)) + } + return nullptr; + } + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr ::cuda::std::uint32_t sizeBytes() const + { + return mAllocatedSize; + } +}; +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/constant_assert.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/constant_assert.cuh new file mode 100644 index 00000000..8d5fe329 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/constant_assert.cuh @@ -0,0 +1,52 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +/* + * _WS_CONSTANT_ASSERT: an assertion that is intended to be verified at compile time. + * + * A _WS_CONSTANT_ASSERT asserts something that the compiler (optimizer) can verify + * at compile time. Therefore, it does not result in an actual call to assert in + * the compiled binary. This allows checking various properties that cannot be + * verified using static_assert. + * + * To ensure that all _WS_CONSTANT_ASSERTs are in fact eliminated, compile with + * -D_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME. With this macro defined, any + * _WS_CONSTANT_ASSERT failure will output illegal PTX containing the error message. + * As a result, compilation will fail. + * + * Compiling with -D_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME has the additional + * advantage that violating any of the assertions can be detected at compile + * time and before even running the code. + * + */ + +#if defined(_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME) && defined(__CUDA_ARCH__) +// When _WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME is defined and compiling for device, output illegal PTX. +// This causes the compilation to fail. +# define _WS_CONSTANT_ASSERT(expr, msg) \ + do \ + { \ + if (!(expr)) \ + { \ + asm volatile(".pragma \"\n" __FILE__ "(" _CCCL_TO_STRING( \ + __LINE__) "): %0" \ + ": error: constant assertion failed with '" msg "'\n\";" ::"C"(__func__)); \ + } \ + } while (0) +#else +// Host or !_WARPSPEED_FORCE_ASSERT_AT_COMPILE_TIME +# define _WS_CONSTANT_ASSERT(expr, msg) _CCCL_ASSERT((expr), msg) +#endif diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/look_ahead.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/look_ahead.cuh new file mode 100644 index 00000000..d9e7cbfe --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/look_ahead.cuh @@ -0,0 +1,338 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !_CCCL_HAS_NV_ATOMIC_BUILTINS() +# include +#endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS() + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +[[nodiscard]] _CCCL_HOST_DEVICE_API _CCCL_CONSTEVAL ::cuda::std::size_t max_native_atomic_size() noexcept +{ +#if _CCCL_CUDA_COMPILER(NVHPC) + return 8; +#else // ^^^ _CCCL_CUDA_COMPILER(NVHPC) ^^^ / vvv !_CCCL_CUDA_COMPILER(NVHPC) vvv + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return 16;), (return 8;)) +#endif // !_CCCL_CUDA_COMPILER(NVHPC) +} + +enum scan_state : ::cuda::std::uint32_t +{ + empty = 0, + tile_aggregate = 1, +}; + +template +struct tile_state_unaligned_t +{ + scan_state state; + AccumT value; +}; + +// some older nvcc versions do not evaluate next_power_of_two() at compile time when called inside an attribute, so we +// have to force constant evaluation by assigning the result to a template parameter +template ))> +struct alignas(_Alignment) tile_state_t : tile_state_unaligned_t +{}; + +#if __cccl_ptx_isa >= 860 + +template +_CCCL_DEVICE_API void +storeTileAggregate(tile_state_t* ptrTileStates, scan_state scanState, AccumT aggr, int index, int num_tiles) +{ + _CCCL_ASSERT(::cuda::is_aligned(ptrTileStates, alignof(tile_state_t)), ""); + _CCCL_ASSERT(index >= 0 && index < num_tiles, "Reading out of bounds tile state"); + + if constexpr (sizeof(tile_state_t) <= cub::detail::warpspeed::max_native_atomic_size() + && ::cuda::is_trivially_copyable_v>) + { + static_assert(::cuda::is_power_of_two(sizeof(tile_state_t))); + tile_state_t tmp{scanState, aggr}; + +# if _CCCL_HAS_NV_ATOMIC_BUILTINS() + __nv_atomic_store(ptrTileStates + index, &tmp, __NV_ATOMIC_RELAXED, __NV_THREAD_SCOPE_DEVICE); +# else // ^^^ _CCCL_HAS_NV_ATOMIC_BUILTINS() ^^^ / vvv !_CCCL_HAS_NV_ATOMIC_BUILTINS() vvv + ::cuda::atomic_ref, ::cuda::std::thread_scope_device>{ptrTileStates[index]}.store( + tmp, ::cuda::std::memory_order_relaxed); +# endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS() + } + else + { + ThreadStore(&ptrTileStates[index].value, aggr); + using state_int = ::cuda::std::underlying_type_t; + store_release(reinterpret_cast(&ptrTileStates[index].state), scanState); + } +} + +template +_CCCL_DEVICE_API tile_state_t loadTileAggregate(tile_state_t* ptrTileStates, int index, int num_tiles) +{ + _CCCL_ASSERT(::cuda::is_aligned(ptrTileStates, alignof(tile_state_t)), ""); + _CCCL_ASSERT(index >= 0 && index < num_tiles, "Reading out of bounds tile state"); + + tile_state_t res; + if constexpr (sizeof(tile_state_t) <= cub::detail::warpspeed::max_native_atomic_size() + && ::cuda::is_trivially_copyable_v>) + { + static_assert(::cuda::is_power_of_two(sizeof(tile_state_t))); +# if _CCCL_HAS_NV_ATOMIC_BUILTINS() + __nv_atomic_load(ptrTileStates + index, &res, __NV_ATOMIC_RELAXED, __NV_THREAD_SCOPE_DEVICE); +# else // ^^^ _CCCL_HAS_NV_ATOMIC_BUILTINS() ^^^ / vvv !_CCCL_HAS_NV_ATOMIC_BUILTINS() vvv + res = ::cuda::atomic_ref, ::cuda::std::thread_scope_device>{ptrTileStates[index]}.load( + ::cuda::std::memory_order_relaxed); +# endif // !_CCCL_HAS_NV_ATOMIC_BUILTINS() + } + else + { + using state_int = ::cuda::std::underlying_type_t; + res.state = static_cast(load_acquire(reinterpret_cast(&ptrTileStates[index].state))); + res.value = ThreadLoad(&ptrTileStates[index].value); + } + return res; +} + +// warpLoadLookahead loads tmp states: +// idxTileCur + [0; 32 * numTileStatesPerThread[ +// +// The states are loaded in laneId order and warp-strided: +// +// outTmpStates[0] contains: +// Lane 0: idxTileCur + 0 +// Lane 1: idxTileCur + 1 +// ... +// Lane 31: idxTileCur + 31 +// +// outTmpStates[1] contains: +// Lane 0: idxTileCur + 32 +// ... +// Lane 31 idxTileCur + 63 +// +// If the index idxTileCur + ii of the loaded state is equal to or exceeds idxTileNext, i.e., idxTileCur + ii >= +// idxTileNext, then the state is not loaded from memory and set to empty. +template +_CCCL_DEVICE_API void warpLoadLookahead( + int laneIdx, + tile_state_t (&outTileStates)[numTileStatesPerThread], + tile_state_t* ptrTileStates, + int idxTileCur, + int idxTileNext, + int num_tiles) +{ + for (int i = 0; i < numTileStatesPerThread; ++i) + { + const int idxTileLookahead = idxTileCur + 32 * i + laneIdx; + if (idxTileLookahead < idxTileNext) + { + outTileStates[i] = loadTileAggregate(ptrTileStates, idxTileLookahead, num_tiles); + } + else + { + // If we are looking ahead of idxTileNext, then set state to empty + outTileStates[i].state = scan_state::empty; + } + } +} + +// warpIncrementalLookahead takes the latest known aggrExclusiveCtaPrev and its tile index, idxTilePrev (which's +// aggregate is NOT included in aggrExclusiveCtaPrev), and computes the aggrExclusiveCta for the next tile of interest, +// idxTileNext (where the returned value will NOT include the aggregate of idxTileNext). +// +// It does so by loading states in chunks of 32 * numTileStatesPerThread elements, starting from idxTilePrev + 1. From +// the chunk of states, it tries to advance its knowledge of aggrExclusiveCta as much as possible. It loops until it can +// calculate the value of aggrExclusiveCta from the preceding states. +// +// The function must be called from a single warp. All passed arguments must be warp-uniform. +template +[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE AccumT warpIncrementalLookahead( + SpecialRegisters specialRegisters, + tile_state_t* ptrTileStates, + const int idxTilePrev, + const AccumT aggrExclusiveCtaPrev, + const int idxTileNext, + ScanOpT& scan_op, + const int num_tiles) +{ + const int laneIdx = static_cast(specialRegisters.laneIdx); + [[maybe_unused]] const ::cuda::std::uint32_t lanemaskEq = ::cuda::ptx::get_sreg_lanemask_eq(); + + int idxTileCur = idxTilePrev; + AccumT aggrExclusiveCtaCur = aggrExclusiveCtaPrev; + + using warp_reduce_t = WarpReduce; + static_assert(::cuda::std::is_same_v>, + "WarpReduce for a full warp must not require temporary storage"); + [[maybe_unused]] typename warp_reduce_t::TempStorage temp_storage; + + while (idxTileCur < idxTileNext) + { + tile_state_t regTmpStates[numTileStatesPerThread]; + warpLoadLookahead(laneIdx, regTmpStates, ptrTileStates, idxTileCur, idxTileNext, num_tiles); + + for (int idx = 0; idx < numTileStatesPerThread; ++idx) + { + // Bitmask with 1 bits indicating which lane has a tile aggregate + const ::cuda::std::uint32_t warp_has_aggregate_mask = + __ballot_sync(0xffffffffu, regTmpStates[idx].state == scan_state::tile_aggregate); + + // Bitmask with 1 bits for all rightmost lanes having a tile aggregate + const ::cuda::std::uint32_t warp_right_aggregates_mask = warp_has_aggregate_mask & (~warp_has_aggregate_mask - 1); + + // Cannot reduce if no rightmost tile aggregates + if (warp_right_aggregates_mask == 0) + { + break; + } + + const ::cuda::std::uint32_t warp_right_aggregates_count = ::cuda::std::popcount(warp_right_aggregates_mask); + + // Accumulate the rightmost tile aggregates + AccumT local_aggr; + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_80, + ({ // NOTE: Inlined from warp_reduce_shfl + if constexpr (is_warp_redux_op_supported_sm80) + { + const bool use_value = lanemaskEq & warp_right_aggregates_mask; + const AccumT value = use_value ? regTmpStates[idx].value : cuda::identity_element(); + local_aggr = cub::detail::warp_redux_sm80(value, ~0, scan_op); + } + else + { + // TODO(bgruber): this generates a LOT of SASS. I think it can do better. + local_aggr = + warp_reduce_t{temp_storage}.Reduce(regTmpStates[idx].value, scan_op, warp_right_aggregates_count); + } + }), + (local_aggr = + warp_reduce_t{temp_storage}.Reduce(regTmpStates[idx].value, scan_op, warp_right_aggregates_count);)) + + // We never initialized aggrExclusiveCtaCur when starting look ahead at tile 0 + aggrExclusiveCtaCur = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr); + idxTileCur += warp_right_aggregates_count; + + // we can only continue on the next 32 tile states, if we consumed all 32 of this iteration + if (warp_right_aggregates_count < 32) + { + break; + } + } + } + + return aggrExclusiveCtaCur; // must only be valid in lane_0 +} + +// Deterministic version of warpIncrementalLookahead that returns the same aggrExclusiveCta. The difference is that it +// always starts the lookahead from a tile index that is a multiple of 32. The left pointer (idxTilePrev) is itself +// always a multiple of 32, as it starts at 0 and is only ever advanced by whole batches of 32, so the lookahead resumes +// from there directly. Because every reduction begins at the same fixed tiles, no matter which tiles happened to finish +// first, the order in which values are summed is always the same and the result is identical on every run. +// idxTilePrev/aggrExclusiveCtaPrev are updated by reference to the last multiple of 32. +template +[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE AccumT warpIncrementalLookaheadStable( + SpecialRegisters specialRegisters, + tile_state_t* ptrTileStates, + int& idxTilePrev, + AccumT& aggrExclusiveCtaPrev, + const int idxTileNext, + ScanOpT& scan_op, + const int num_tiles) +{ + const int laneIdx = static_cast(specialRegisters.laneIdx); + const ::cuda::std::uint32_t lanemaskEq = ::cuda::ptx::get_sreg_lanemask_eq(); + + int idxTileCur = idxTilePrev; + AccumT aggrExclusiveCtaCur = aggrExclusiveCtaPrev; + + using warp_reduce_t = WarpReduce; + static_assert(::cuda::std::is_same_v>, + "WarpReduce for a full warp must not require temporary storage"); + [[maybe_unused]] typename warp_reduce_t::TempStorage temp_storage; + + while (idxTileCur < idxTileNext) + { + tile_state_t regTmpStates[numTileStatesPerThread]; + warpLoadLookahead(laneIdx, regTmpStates, ptrTileStates, idxTileCur, idxTileNext, num_tiles); + + for (int idx = 0; idx < numTileStatesPerThread; ++idx) + { + // Bitmask with 1 bits indicating which lane has a tile aggregate + const ::cuda::std::uint32_t warp_has_aggregate_mask = + __ballot_sync(0xffffffffu, regTmpStates[idx].state == scan_state::tile_aggregate); + + // Bitmask with 1 bits for the contiguous run of lanes having a tile aggregate starting from LSB + const ::cuda::std::uint32_t warp_right_aggregates_mask = warp_has_aggregate_mask & (~warp_has_aggregate_mask - 1); + + const ::cuda::std::uint32_t warp_right_aggregates_count = ::cuda::std::popcount(warp_right_aggregates_mask); + + // Only reduce once 32 contiguous tile aggregates are available, so the reduction order is fixed. + const ::cuda::std::uint32_t expected_count = + static_cast<::cuda::std::uint32_t>(::cuda::std::min(32, idxTileNext - idxTileCur)); + if (warp_right_aggregates_count < expected_count) + { + break; + } + + const bool use_value = lanemaskEq & warp_right_aggregates_mask; + const AccumT value = use_value ? regTmpStates[idx].value : cuda::identity_element(); + const AccumT local_aggr = warp_reduce_t{temp_storage}.Reduce(value, scan_op); + + if (expected_count == 32) + { + aggrExclusiveCtaCur = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr); + idxTileCur += 32; + } + else + { + const AccumT full_aggr = idxTileCur == 0 ? local_aggr : scan_op(aggrExclusiveCtaCur, local_aggr); + idxTilePrev = idxTileCur; + aggrExclusiveCtaPrev = aggrExclusiveCtaCur; + return full_aggr; + } + } + } + + // Only reached when idxTileNext is a multiple of 32; otherwise the final partial batch full aggregate returns inside + // the loop above. + idxTilePrev = idxTileNext; + aggrExclusiveCtaPrev = aggrExclusiveCtaCur; + return aggrExclusiveCtaCur; // must only be valid in lane_0 +} + +#endif // __cccl_ptx_isa >= 860 +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/make_warp_uniform.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/make_warp_uniform.cuh new file mode 100644 index 00000000..1462bf0c --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/make_warp_uniform.cuh @@ -0,0 +1,42 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +// Move register to uniform register + +// For int32_t and uint32_t, we can use the CREDUX instruction, which is coupled and has a constant latency. +// For 64-bit types, we still use __shfl_sync + +[[nodiscard]] _CCCL_DEVICE_API inline int makeWarpUniform(int x) +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return __reduce_min_sync(~0, x);), (return x;)); +} + +[[nodiscard]] _CCCL_DEVICE_API inline ::cuda::std::uint32_t makeWarpUniform(::cuda::std::uint32_t x) +{ + NV_IF_ELSE_TARGET(NV_PROVIDES_SM_90, (return __reduce_min_sync(~0, x);), (return x;)); +} + +[[nodiscard]] _CCCL_DEVICE_API inline ::cuda::std::uint64_t makeWarpUniform(::cuda::std::uint64_t x) +{ + return __shfl_sync(~0, x, 0); +} +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/optimize_smem_ptr.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/optimize_smem_ptr.cuh new file mode 100644 index 00000000..80ed6e90 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/optimize_smem_ptr.cuh @@ -0,0 +1,38 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +template +[[nodiscard]] _CCCL_DEVICE_API _Tp* optimizeSmemPtr(const _Tp* smemGeneric) +{ + // See https://nvbugspro.nvidia.com/bug/4907996 + + // 1. Convert to 32-bit shared memory pointer + ::cuda::std::uint32_t smem32 = __cvta_generic_to_shared(smemGeneric); + // 2. Pretend to NVVM that the 32-bit pointer is modified. This is required to avoid NVVM constant + // propagation from pulling the smem32 definition into loops and branches in subsequent code. + asm("" : "+r"(smem32)); + // 3. Make a generic pointer to smem that is constructed using `__cvta_shared_to_generic`. This + // benefits from an + // optimization pass in NVVM that performs the following simplification: + // __cvta_generic_to_shared(__cvta_shared_to_generic(x)) => x. + // In our case, `x` is smem32, which is exactly what we want. + return reinterpret_cast<_Tp*>(__cvta_shared_to_generic(smem32)); +} +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_phase.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_phase.cuh new file mode 100644 index 00000000..dc2e12ba --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_phase.cuh @@ -0,0 +1,45 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +template +struct SmemPhase +{ + SmemResourceRaw& mSmemResourceRaw; + int mCurPhase; + + _CCCL_DEVICE_API SmemPhase(SmemResourceRaw& smemResourceRaw, int phase) noexcept + : mSmemResourceRaw(smemResourceRaw) + , mCurPhase(phase) + {} + + [[nodiscard]] _CCCL_DEVICE_API SmemRef<_Tp> acquireRef() + { + // Wait on barrier + mSmemResourceRaw.acquire(mCurPhase); + // Return ref + return SmemRef<_Tp>(mSmemResourceRaw, mCurPhase); + } +}; +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_ref.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_ref.cuh new file mode 100644 index 00000000..e092a7cb --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_ref.cuh @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +template +struct SmemRef +{ + SmemResourceRaw& mSmemResourceRaw; + int mCurPhase; + bool mTxCountIsSet = false; + int mTxCount = 0; + bool mDoFenceLdsToAsyncProxy = false; + + _CCCL_DEVICE_API SmemRef(SmemResourceRaw& smemResourceRaw, int phase) noexcept + : mSmemResourceRaw(smemResourceRaw) + , mCurPhase(phase) + {} + // SmemRef is a non-copyable, non-movable type. It must be passed by (mutable) + // reference to be useful. The reason is that it in case of an accidental copy + // or move the destructor is called twice. This leads to double-arrivals on + // barriers and results in deadlock or a hardware fault. + SmemRef(const SmemRef&) = delete; // Delete copy constructor + SmemRef(SmemRef&&) = delete; // Delete move constructor + SmemRef& operator=(const SmemRef&) = delete; // Delete copy assignment + SmemRef& operator=(SmemRef&&) = delete; // Delete move assignment + + _CCCL_DEVICE_API ~SmemRef() + { + if (mDoFenceLdsToAsyncProxy) + { + mSmemResourceRaw.fenceLdsToAsyncProxy(); + } + if (mTxCountIsSet) + { + mSmemResourceRaw.releaseTx(mCurPhase, mTxCount); + } + else + { + mSmemResourceRaw.release(mCurPhase); + } + } + + [[nodiscard]] _CCCL_DEVICE_API _Tp& data() noexcept + { + return *static_cast<_Tp*>(mSmemResourceRaw.data()); + } + + [[nodiscard]] _CCCL_DEVICE_API int sizeBytes() const noexcept + { + return mSmemResourceRaw.mSizeBytes; + } + + [[nodiscard]] _CCCL_DEVICE_API uint64_t* ptrCurBarrierRelease() + { + return mSmemResourceRaw.ptrCurBarrierRelease(mCurPhase); + } + + _CCCL_DEVICE_API void squadIncreaseTxCount(const Squad& squad, int txCount) + { + mTxCountIsSet = true; + // Only leader thread increments txCount + txCount = squad.isLeaderThread() ? txCount : 0; + mTxCount += txCount; + } + + _CCCL_DEVICE_API void setFenceLdsToAsyncProxy() noexcept + { + mDoFenceLdsToAsyncProxy = true; + } +}; +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_resource.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_resource.cuh new file mode 100644 index 00000000..93ba6c48 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_resource.cuh @@ -0,0 +1,60 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +template +struct SmemResource : SmemResourceRaw +{ + template + _CCCL_HOST_DEVICE_API SmemResource(SyncHandler& syncHandler, _Tp (&smemBuffer)[stageCount]) + : SmemResourceRaw(syncHandler, smemBuffer, sizeof(smemBuffer[0]), sizeof(smemBuffer[0]), stageCount) + {} + + _CCCL_HOST_DEVICE_API constexpr SmemResource( + SyncHandler& syncHandler, SmemAllocator& smemAllocator, Stages stages, Elems elems = Elems{1}) + : SmemResourceRaw(makeSmemResourceRaw(syncHandler, smemAllocator, stages, elems)) + {} + + [[nodiscard]] _CCCL_DEVICE_API SmemStage<_Tp> nextStage() noexcept + { + return SmemStage<_Tp>(*this); + } + +private: + [[nodiscard]] _CCCL_HOST_DEVICE_API static constexpr SmemResourceRaw + makeSmemResourceRaw(SyncHandler& syncHandler, SmemAllocator& smemAllocator, Stages stages, Elems elems = Elems{1}) + { + int align = alignof(_Tp); + int sizeBytes = ::cuda::std::to_underlying(elems) * sizeof(_Tp); + int strideBytes = sizeBytes; + + void* ptrBase = smemAllocator.alloc(::cuda::std::to_underlying(stages) * strideBytes, align); + return {syncHandler, ptrBase, sizeBytes, strideBytes, ::cuda::std::to_underlying(stages)}; + } +}; +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_resource_raw.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_resource_raw.cuh new file mode 100644 index 00000000..938a9608 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_resource_raw.cuh @@ -0,0 +1,182 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +struct SmemResourceRaw +{ + static constexpr int mMaxNumPhases = 4; + + int mStageCurrent = 0; + + int mResourceHandle; + ::cuda::std::uint8_t* mPtrBase{}; + int mSizeBytes; + int mStride; + int mStageCount; + int mNumPhases = 0; + + ::cuda::std::uint64_t* mPtrBar[mMaxNumPhases]{}; + int mParity[mMaxNumPhases]{}; + + _CCCL_HOST_DEVICE_API constexpr SmemResourceRaw( + SyncHandler& syncHandler, void* ptrBase, int sizeBytes, int strideBytes, int stageCount) noexcept + : mResourceHandle(syncHandler.registerResource(stageCount)) + , mSizeBytes(sizeBytes) + , mStride(strideBytes) + , mStageCount(stageCount) + { + // we don't need the pointer during constant evaluation (and casting is not allowed) + if (!::cuda::std::is_constant_evaluated()) + { + mPtrBase = static_cast<::cuda::std::uint8_t*>(ptrBase); + } + + for (int pi = 0; pi < mMaxNumPhases; ++pi) + { + mParity[pi] = pi == 0 ? 1 : 0; + } + } + + template + _CCCL_HOST_DEVICE_API constexpr void + addPhase(SyncHandler& syncHandler, ::cuda::std::uint64_t* ptrBarrier, const SquadDesc (&squads)[numSquads]) + { + int numOwningThreads = squadCountThreads(squads); + + int curPhase = mNumPhases; + mNumPhases++; + + syncHandler.registerPhase(mResourceHandle, numOwningThreads, ptrBarrier); + mPtrBar[curPhase] = ptrBarrier; + } + + template + _CCCL_HOST_DEVICE_API constexpr void + addPhase(SyncHandler& syncHandler, SmemAllocator& smemAllocator, const SquadDesc (&squads)[numSquads]) + { + void* ptrBar_raw = smemAllocator.alloc(mStageCount * sizeof(::cuda::std::uint64_t), alignof(::cuda::std::uint64_t)); + // we don't need the pointer during constant evaluation (and casting is not allowed) + ::cuda::std::uint64_t* ptrBar = nullptr; + if (!::cuda::std::is_constant_evaluated()) + { + ptrBar = static_cast<::cuda::std::uint64_t*>(ptrBar_raw); + } + addPhase(syncHandler, ptrBar, squads); + } + + _CCCL_HOST_DEVICE_API void + addPhase(SyncHandler& syncHandler, ::cuda::std::uint64_t* ptrBarrier, const SquadDesc& squad) + { + const SquadDesc squads[] = {squad}; + addPhase(syncHandler, ptrBarrier, squads); + } + + _CCCL_HOST_DEVICE_API constexpr void + addPhase(SyncHandler& syncHandler, SmemAllocator& smemAllocator, const SquadDesc& squad) + { + const SquadDesc squads[] = {squad}; + addPhase(syncHandler, smemAllocator, squads); + } + + _CCCL_DEVICE_API void incrementStage() + { + if (mStageCurrent == mStageCount - 1) + { + mStageCurrent = 0; + // We loop over all phases with a conditional on resNumPhases. If we + // directly loop over only resNumPhases, then the SROA optimization does + // not kick in and the mParity array is spilled to the stack. + for (int pi = 0; pi < mMaxNumPhases; ++pi) + { + if (pi < mNumPhases) + { + mParity[pi] ^= 1; + } + } + } + else + { + mStageCurrent++; + } + } + + _CCCL_DEVICE_API void* data() + { + return (void*) (mPtrBase + mStageCurrent * mStride); + } + + [[nodiscard]] _CCCL_DEVICE_API ::cuda::std::uint64_t* ptrCurBarrierRelease(int phase) + { + ::cuda::std::uint64_t* ptrBarPhase = mPtrBar[phase]; + _WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit."); + return &ptrBarPhase[mStageCurrent]; + } + _CCCL_DEVICE_API void release(int phase) + { + _WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit."); + ::cuda::ptx::mbarrier_arrive(ptrCurBarrierRelease(phase)); + } + + _CCCL_DEVICE_API void releaseTx(int phase, int txCount) + { + _WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit."); + ::cuda::ptx::mbarrier_arrive_expect_tx( + ::cuda::ptx::sem_release, ::cuda::ptx::scope_cta, ::cuda::ptx::space_shared, ptrCurBarrierRelease(phase), txCount); + } + + _CCCL_DEVICE_API void fenceLdsToAsyncProxy() + { + ::cuda::ptx::fence_proxy_async(::cuda::ptx::space_shared); + } + + _CCCL_DEVICE_API void releaseLdsToAsyncProxy(int phase) + { + // First fence + fenceLdsToAsyncProxy(); + // Then perform a normal release + release(phase); + } + + _CCCL_DEVICE_API void acquire(int phase) + { + _WS_CONSTANT_ASSERT(phase < mNumPhases, "Phase exceeds limit."); + + // The release of the previous phase occurs on the `phase - 1` barrier. So + // that is what we wait on. + int phaseAcq = (mNumPhases + phase - 1) % mNumPhases; + ::cuda::std::uint64_t* ptrBarPhase = mPtrBar[phaseAcq]; + + while (!::cuda::ptx::mbarrier_try_wait_parity(&ptrBarPhase[mStageCurrent], mParity[phase])) + { + } + } +}; +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_stage.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_stage.cuh new file mode 100644 index 00000000..d74d7caf --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/resource/smem_stage.cuh @@ -0,0 +1,91 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +template +struct SmemStage +{ + SmemResourceRaw& mSmemResourceRaw; + + _CCCL_DEVICE_API SmemStage(SmemResourceRaw& smemResourceRaw) noexcept + : mSmemResourceRaw(smemResourceRaw) + {} + + _CCCL_DEVICE_API ~SmemStage() + { + mSmemResourceRaw.incrementStage(); + } + + // SmemStage is a non-copyable, non-movable type. It must be passed by (mutable) + // reference to be useful. The reason is that it in case of an accidental copy + // or move the destructor is called twice. This leads to double-increment of + // the stage index and results in deadlock or a hardware fault. + SmemStage(const SmemStage&) = delete; // Delete copy constructor + SmemStage(SmemStage&&) = delete; // Delete move constructor + SmemStage& operator=(const SmemStage&) = delete; // Delete copy assignment + SmemStage& operator=(const SmemStage&&) = delete; // Delete move assignment +}; + +// Helper: Container to expose SmemPhase for structured binding +template +struct SmemPhaseStructuredBinding +{ + SmemResourceRaw& mSmemResourceRaw; + + template <::cuda::std::size_t _Index> + [[nodiscard]] _CCCL_DEVICE_API SmemPhase<_Tp> get() const + { + return SmemPhase<_Tp>(mSmemResourceRaw, _Index); + } +}; + +// The binding function +template <::cuda::std::size_t numPhases, typename _Tp> +[[nodiscard]] _CCCL_DEVICE_API SmemPhaseStructuredBinding<_Tp, numPhases> bindPhases(SmemStage<_Tp>& smemStage) +{ + _WS_CONSTANT_ASSERT(smemStage.mSmemResourceRaw.mNumPhases == numPhases, + "Number of bound phases must match resource phases."); + + return SmemPhaseStructuredBinding<_Tp, numPhases>{smemStage.mSmemResourceRaw}; +} +} // namespace detail::warpspeed + +CUB_NAMESPACE_END + +// Tuple protocol specializations +namespace std +{ +template +struct tuple_size> +{ + static constexpr size_t value = numPhases; +}; + +template +struct tuple_element<_Index, CUB_NS_QUALIFIER::detail::warpspeed::SmemPhaseStructuredBinding<_Tp, numPhases>> +{ + using type = CUB_NS_QUALIFIER::detail::warpspeed::SmemPhase<_Tp>; +}; +} // namespace std diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/special_registers.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/special_registers.cuh new file mode 100644 index 00000000..e18f34bd --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/special_registers.cuh @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +// Commonly used special registers that we should cache in registers or uniform +// registers. +struct SpecialRegisters +{ + const ::cuda::std::uint32_t clusterCtaRank; + const ::cuda::std::uint32_t blockIdxX; + const ::cuda::std::uint32_t threadIdxX; + const ::cuda::std::uint32_t warpIdx; + const ::cuda::std::uint32_t laneIdx; +}; + +[[nodiscard]] _CCCL_DEVICE_API inline SpecialRegisters getSpecialRegisters() +{ + ::cuda::std::uint32_t clusterCtaRank = ::cuda::ptx::get_sreg_cluster_ctarank(); + ::cuda::std::uint32_t threadIdxX = threadIdx.x; + ::cuda::std::uint32_t warpIdx = makeWarpUniform(threadIdxX / 32); + return {clusterCtaRank, blockIdx.x, threadIdxX, warpIdx, ::cuda::ptx::get_sreg_laneid()}; +} +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/load_store.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/load_store.cuh new file mode 100644 index 00000000..104646c2 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/load_store.cuh @@ -0,0 +1,388 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +#if __cccl_ptx_isa >= 860 + +template +struct CpAsyncOobInfo +{ + // The aligned up and down pointers below must be ::cuda::std::byte*, since the nearest aligned up/down ptr may not + // point to a multiple of sizeof(Tp). E.g. a uchar3* pointing to address 0x...5 will be aligned down to 0x...0, but + // that's not a valid start for an uchar3 in that array. So we must express all aligned pointers in bytes here. + + ::cuda::std::byte* ptrGmem; + ::cuda::std::byte* ptrGmemStartAlignDown; + ::cuda::std::byte* ptrGmemStartAlignUp; + ::cuda::std::byte* ptrGmemEnd; + ::cuda::std::byte* ptrGmemEndAlignDown; + ::cuda::std::byte* ptrGmemEndAlignUp; + ::cuda::std::uint32_t overCopySizeBytes; + ::cuda::std::uint32_t underCopySizeBytes; + ::cuda::std::uint32_t origCopySizeBytes; + ::cuda::std::uint32_t smemStartSkipBytes; // ptrSmem + smemStartSkipBytes will point to the first valid element copied + // from ptrGmem + ::cuda::std::uint32_t smemEndBytesAfter16BBoundary; // number of bytes after the last 16B boundary in GMEM/SMEM that + // still contains valid (partial) elements +}; + +template +_CCCL_DEVICE_API _CCCL_FORCEINLINE CpAsyncOobInfo prepareCpAsyncOob(Tp* ptrGmem, ::cuda::std::uint32_t sizeElem) +{ + auto ptrGmemBytes = reinterpret_cast<::cuda::std::byte*>(ptrGmem); + auto ptrGmemEnd = reinterpret_cast<::cuda::std::byte*>(ptrGmem + sizeElem); + + // We will copy from [ptrGmemBase, ptrGmemEnd). Both pointers have to be 16B aligned. + ::cuda::std::byte* ptrGmemStartAlignDown = ::cuda::align_down(ptrGmemBytes, 16); + ::cuda::std::byte* ptrGmemStartAlignUp = ::cuda::align_up(ptrGmemBytes, 16); + ::cuda::std::byte* ptrGmemEndAlignUp = ::cuda::align_up(ptrGmemEnd, 16); + ::cuda::std::byte* ptrGmemEndAlignDown = ::cuda::align_down(ptrGmemEnd, 16); + + // Compute the final copy size in bytes. It can be either sizeElem or sizeElem + 16 / sizeof(T). + const auto origCopySizeBytes = static_cast<::cuda::std::uint32_t>(sizeof(Tp) * sizeElem); + const auto overCopySizeBytes = static_cast<::cuda::std::uint32_t>(ptrGmemEndAlignUp - ptrGmemStartAlignDown); + auto underCopySizeBytes = static_cast<::cuda::std::uint32_t>(ptrGmemEndAlignDown - ptrGmemStartAlignUp); + if (origCopySizeBytes < underCopySizeBytes) + { + // If ptrGmemStart and ptrGmemEnd are aligned to [1, .., 15] bytes, then + // when we align the one up and the other down we get overflow. We check for + // that here. In that case, the undercopy size is zero. + underCopySizeBytes = 0; + } + + _CCCL_ASSERT(overCopySizeBytes % 16 == 0, ""); + _CCCL_ASSERT(underCopySizeBytes % 16 == 0, ""); + + return { + ptrGmemBytes, + ptrGmemStartAlignDown, + ptrGmemStartAlignUp, + ptrGmemEnd, + ptrGmemEndAlignDown, + ptrGmemEndAlignUp, + overCopySizeBytes, + underCopySizeBytes, + origCopySizeBytes, + static_cast<::cuda::std::uint32_t>(ptrGmemBytes - ptrGmemStartAlignDown), + static_cast<::cuda::std::uint32_t>(ptrGmemEnd - ptrGmemEndAlignDown), + }; +} + +template +_CCCL_DEVICE_API void squadLoadBulk(Squad squad, SmemRef& refDestSmem, CpAsyncOobInfo cpAsyncOobInfo) +{ + ::cuda::std::byte* ptrSmem = refDestSmem.data().inout; + _CCCL_ASSERT(::cuda::is_aligned(ptrSmem, 16), ""); + ::cuda::std::uint64_t* ptrBar = refDestSmem.ptrCurBarrierRelease(); + + if constexpr (alignof(Tp) >= 16) + { + // for alignments larger than 16, we can just bulk copy, even just a single element + if (squad.isLeaderThread()) + { + ::cuda::ptx::cp_async_bulk( + ::cuda::std::conditional_t<__cccl_ptx_isa >= 860, ::cuda::ptx::space_shared_t, ::cuda::ptx::space_cluster_t>{}, + ::cuda::ptx::space_global, + ptrSmem, + cpAsyncOobInfo.ptrGmem, + cpAsyncOobInfo.origCopySizeBytes, + ptrBar); + } + refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.underCopySizeBytes); + } + else + { + // for alignments smaller than 16, we can overcopy but need to declare the ignored bytes left and right +# if __cccl_ptx_isa >= 920 + if (squad.isLeaderThread()) + { + ::cuda::ptx::cp_async_bulk_ignore_oob( + ::cuda::ptx::space_shared, + ::cuda::ptx::space_global, + ptrSmem, + cpAsyncOobInfo.ptrGmemStartAlignDown, + cpAsyncOobInfo.overCopySizeBytes, + /* ignore left */ cpAsyncOobInfo.smemStartSkipBytes, + /* ignore right */ cpAsyncOobInfo.ptrGmemEndAlignUp - cpAsyncOobInfo.ptrGmemEnd, + ptrBar); + } + refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.overCopySizeBytes); +# else // __cccl_ptx_isa >= 920 + // if we don't have cp_async_bulk_ignore_oob, we have to undercopy and copy head and tail elements manually + + // handle small copies first. If we have less than 16 bytes we may not straddle a 16B boundary + if (cpAsyncOobInfo.origCopySizeBytes < 16) + { + const auto elemCount = cpAsyncOobInfo.origCopySizeBytes / sizeof(Tp); + _CCCL_ASSERT(elemCount <= squad.threadCount(), ""); + if (squad.threadRank() < elemCount) + { + reinterpret_cast(ptrSmem + cpAsyncOobInfo.smemStartSkipBytes)[squad.threadRank()] = + reinterpret_cast(cpAsyncOobInfo.ptrGmem)[squad.threadRank()]; + } + return; // no bulk copy has been performed so we don't need to update the tx count of any barrier + } + + // copies larger than 16 byte which straddle at least one 16B boundary, so we have dedicated start and end copies + + const bool doStartCopy = cpAsyncOobInfo.smemStartSkipBytes > 0; + + ::cuda::std::byte* ptrSmemMiddle = ptrSmem; + if (doStartCopy) + { + ptrSmemMiddle += 16; + } + + // TODO(bgruber): we could skip the middle if underCopySizeBytes is zero + if (squad.isLeaderThread()) + { + ::cuda::ptx::cp_async_bulk( + ::cuda::std::conditional_t<__cccl_ptx_isa >= 860, ::cuda::ptx::space_shared_t, ::cuda::ptx::space_cluster_t>{}, + ::cuda::ptx::space_global, + ptrSmemMiddle, + cpAsyncOobInfo.ptrGmemStartAlignUp, + cpAsyncOobInfo.underCopySizeBytes, + ptrBar); + } + refDestSmem.squadIncreaseTxCount(squad, cpAsyncOobInfo.underCopySizeBytes); + + // we cannot use Tp to load the head and tail elements, because sizeof(Tp) may be larger than alignof(Tp) + using load_word_t = ::cuda::std::__make_nbit_uint_t; + + const int head_elements = (cpAsyncOobInfo.ptrGmemStartAlignUp - cpAsyncOobInfo.ptrGmem) / sizeof(load_word_t); + const int tail_elements = (cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemEndAlignDown) / sizeof(load_word_t); + _CCCL_ASSERT(head_elements <= squad.threadCount(), ""); + _CCCL_ASSERT(tail_elements <= squad.threadCount(), ""); + load_word_t head_value, tail_value; + if (squad.threadRank() < head_elements) + { + head_value = reinterpret_cast(cpAsyncOobInfo.ptrGmem)[squad.threadRank()]; + } + if (squad.threadRank() < tail_elements) + { + tail_value = reinterpret_cast(cpAsyncOobInfo.ptrGmemEndAlignDown)[squad.threadRank()]; + } + + if (squad.threadRank() < head_elements) + { + reinterpret_cast(ptrSmem + cpAsyncOobInfo.smemStartSkipBytes)[squad.threadRank()] = head_value; + } + if (squad.threadRank() < tail_elements) + { + reinterpret_cast(ptrSmemMiddle + cpAsyncOobInfo.underCopySizeBytes)[squad.threadRank()] = + tail_value; + } +# endif // __cccl_ptx_isa >= 920 + } +} + +_CCCL_DEVICE_API _CCCL_FORCEINLINE void squadStoreMasked16B( + Squad squad, + ::cuda::std::byte* dstGmem, + const ::cuda::std::byte* srcSmem, + ::cuda::std::uint16_t byteMask, + int firstByte, + int lastByte) +{ + NV_IF_ELSE_TARGET( + NV_PROVIDES_SM_100, + (if (::cuda::ptx::elect_sync(~0)) { + ::cuda::ptx::cp_async_bulk_cp_mask( + ::cuda::ptx::space_global, ::cuda::ptx::space_shared, dstGmem, srcSmem, /*size*/ 16, byteMask); + }), + ({ + const int rank = squad.threadRank(); + if (firstByte <= rank && rank < lastByte) + { + dstGmem[rank] = srcSmem[rank]; + } + })); +} + +template +_CCCL_DEVICE_API void +squadStoreBulkSync(Squad squad, CpAsyncOobInfo cpAsyncOobInfo, const ::cuda::std::byte* srcSmem) +{ + // This function performs either 1 copy, or three copies, depending on the + // size and alignment of the output tile in global memory. + // + // If the output tile is contained in a single 16-byte aligned and sized region, then we + // only perform a single masked copy. + // + // If the output tile is larger than 16 bytes or straddles two 16-byte aligned and sized regions, then + // we perform up to three copies: + // - One copy for the first up to 15 bytes at the start of the region. + // - One copy that starts at a 16-byte aligned address and ends at the latest 16-byte aligned address. + // - One copy that cleans up the last up to 15 bytes. + if (squad.isLeaderWarp()) + { + // Acquire shared memory in async proxy + // Perform fence.proxy.async with full warp to avoid BSSY+BSYNC + ::cuda::ptx::fence_proxy_async(::cuda::ptx::space_shared); + +# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3) + // for some reason the optimizer propagates some information from the computation of + // overCopySizeBytes to the masked bulk copy below and generates an unaligned access error. + // The artificial read modification of overCopySizeBytes prevents the propagation here works around this. + // It also solves the issue described in nvbug 5848313 by accident on nvcc 13.2+ + asm volatile("" : "+r"(cpAsyncOobInfo.overCopySizeBytes)); +# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3) + + const bool doStartCopy = cpAsyncOobInfo.smemStartSkipBytes > 0; + const bool doEndCopy = cpAsyncOobInfo.smemEndBytesAfter16BBoundary > 0; + const bool doMiddleCopy = cpAsyncOobInfo.ptrGmemStartAlignUp != cpAsyncOobInfo.ptrGmemEndAlignUp; + + constexpr ::cuda::std::uint16_t byteMask = 0xFFFF; + const ::cuda::std::uint16_t byteMaskStart = byteMask << cpAsyncOobInfo.smemStartSkipBytes; + const ::cuda::std::uint16_t byteMaskEnd = byteMask >> (16 - cpAsyncOobInfo.smemEndBytesAfter16BBoundary) % 16; + // byteMaskStart contains zeroes at the left +# if _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2) + const ::cuda::std::uint16_t byteMaskSmall = byteMaskStart & byteMaskEnd; +# else // _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2) + // `ptxas fatal : (C7907) Internal compiler error`, see nvbug 5848313 + const ::cuda::std::uint16_t byteMaskSmall = + byteMaskStart & (byteMask >> (16 - (cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemStartAlignDown))); +# endif // _CCCL_CUDA_COMPILER(NVCC, >=, 13, 2) + + const ::cuda::std::byte* ptrSmemMiddle = srcSmem; + if (doStartCopy) + { + ptrSmemMiddle += 16; + } + + if (doMiddleCopy) + { + // Copy the middle part. Starting at byte 0 or 16 in shared memory. This + // is the large copy. We perform this one first, so that the compiler can + // (hopefully) hide all the arithmetic behind this instruction. + if (::cuda::ptx::elect_sync(~0)) + { + // need to work around another optimizer bug, see: https://github.com/NVIDIA/cccl/issues/8644 +# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3) + asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemStartAlignUp)); +# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3) + ::cuda::ptx::cp_async_bulk( + ::cuda::ptx::space_global, + ::cuda::ptx::space_shared, + cpAsyncOobInfo.ptrGmemStartAlignUp, + ptrSmemMiddle, + cpAsyncOobInfo.underCopySizeBytes); + } + if (doStartCopy) + { + // need to work around yet another optimizer bug, see: https://github.com/NVIDIA/cccl/issues/8838 +# if _CCCL_CUDA_COMPILER(NVCC, <, 13, 3) + asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemStartAlignDown)); + asm volatile("" : "+l"(srcSmem)); +# endif // _CCCL_CUDA_COMPILER(NVCC, <, 13, 3) + // Copy a subset of the first 16 bytes + squadStoreMasked16B( + squad, + cpAsyncOobInfo.ptrGmemStartAlignDown, + srcSmem, + byteMaskStart, + static_cast(cpAsyncOobInfo.smemStartSkipBytes), + 16); + } + if (doEndCopy) + { +# if _CCCL_CUDA_COMPILER(NVHPC) + // nvc++ seems to have an optimizer bug, crashing with an unaligned access error below. The addresses are fine + // when printed, so let's shake the optimizer a bit. + asm volatile("" : "+l"(cpAsyncOobInfo.ptrGmemEndAlignDown)); +# endif // _CCCL_CUDA_COMPILER(NVHPC) + + // Copy a subset of the last 16 bytes + squadStoreMasked16B( + squad, + cpAsyncOobInfo.ptrGmemEndAlignDown, + ptrSmemMiddle + cpAsyncOobInfo.underCopySizeBytes, + byteMaskEnd, + 0, + static_cast(cpAsyncOobInfo.smemEndBytesAfter16BBoundary)); + } + } + else + { + // Copy a subset of the first 16 bytes + squadStoreMasked16B( + squad, + cpAsyncOobInfo.ptrGmemStartAlignDown, + srcSmem, + byteMaskSmall, + static_cast(cpAsyncOobInfo.smemStartSkipBytes), + static_cast(cpAsyncOobInfo.ptrGmemEnd - cpAsyncOobInfo.ptrGmemStartAlignDown)); + } + // Commit and wait for store to have completed reading from shared memory + ::cuda::ptx::cp_async_bulk_commit_group(); + ::cuda::ptx::cp_async_bulk_wait_group_read(::cuda::ptx::n32_t<0>{}); + } +} + +#endif // __cccl_ptx_isa >= 860 + +template +_CCCL_DEVICE_API void squadLoadSmem(Squad squad, AccumT (&outReg)[ElemPerThread], const InputT* smemBuf) +{ + for (int i = 0; i < ElemPerThread; ++i) + { + const int elem_idx = squad.threadRank() * ElemPerThread + i; + outReg[i] = smemBuf[elem_idx]; + } +} + +template +_CCCL_DEVICE_API void squadStoreSmem(Squad squad, OutputT* smemBuf, const AccumT (&inReg)[ElemPerThread]) +{ + for (int i = 0; i < ElemPerThread; ++i) + { + const int elem_idx = squad.threadRank() * ElemPerThread + i; + smemBuf[elem_idx] = inReg[i]; + } +} + +template +_CCCL_DEVICE_API void +squadStoreSmemPartial(Squad squad, OutputT* smemBuf, const AccumT (&inReg)[ElemPerThread], int beginIndex, int endIndex) +{ + for (int i = 0; i < ElemPerThread; ++i) + { + const int elem_idx = squad.threadRank() * ElemPerThread + i; + if (beginIndex <= elem_idx && elem_idx < endIndex) + { + smemBuf[elem_idx - beginIndex] = inReg[i]; + } + } +} +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/squad.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/squad.cuh new file mode 100644 index 00000000..97a90107 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/squad.cuh @@ -0,0 +1,159 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +// Squad - device squad instance +// +// A squad is a collection of warps that work together in a warp-specialized +// kernel. A warp-specialized kernel has multiple squads that perform part of +// the computation. +// +// The Squad class is a device runtime instance of a squad. It provides +// functionality to determine the rank of the current thread or warp in the +// squad, and to sync all threads in the squad. +struct Squad : SquadDesc +{ + SpecialRegisters mSpecialRegisters; + bool mIsWarpLeader = false; + bool mIsLeaderWarp = false; + + _CCCL_DEVICE_API Squad(SquadDesc squadStatic, SpecialRegisters specialRegisters) + : SquadDesc(squadStatic) + , mSpecialRegisters(specialRegisters) + { + mIsWarpLeader = ::cuda::ptx::elect_sync(~0); + mIsLeaderWarp = warpRank() == 0; + } + + [[nodiscard]] _CCCL_DEVICE_API int warpRank() const + { + return static_cast(mSpecialRegisters.warpIdx % this->warpCount()); + } + + [[nodiscard]] _CCCL_DEVICE_API int threadRank() const + { + return static_cast(mSpecialRegisters.threadIdxX % this->threadCount()); + } + + [[nodiscard]] _CCCL_DEVICE_API bool isLeaderThread() const + { + return mIsWarpLeader && mIsLeaderWarp; + } + + [[nodiscard]] _CCCL_DEVICE_API bool isLeaderWarp() const + { + return mIsLeaderWarp; + } + + [[nodiscard]] _CCCL_DEVICE_API bool isLeaderThreadOfWarp() const + { + return mIsWarpLeader; + } + + _CCCL_DEVICE_API void syncThreads() const + { + // barrier 0 is reserved for __syncthreads(). We use barrier ids 1, ... + const int barrierIdx = this->mSquadIdx + 1; + + __barrier_sync_count(barrierIdx, this->threadCount()); + } +}; +// squadDispatch +// +// squadDispatch is used at the start of the kernel. It takes an array of squad +// descriptors and determines which squad the current thread belongs to. The +// lambda `f: (Squad) -> void` is called with the squad currently active on this +// thread. +// +// Typically, the user will call the kernel body with the active squad. +// +// Implementation notes: +// +// Dispatch to squad based on warp index using a binary search. This balances +// the number of BRA instructions per squad and avoids NVVM inserting BRX +// instructions. BRX instructions require a jump table that is loaded from GCC, +// which incurs latency. +// +// The benefit of this function for fastScan is that adding a new squad doesn't +// require code changes in the dispatch. For low-latency inference, I hope that +// the avoidance of linear search and BRX instructions translates into latency +// reductions. +// +template +_CCCL_DEVICE_API _CCCL_FORCEINLINE void +squadDispatch(SpecialRegisters sr, const SquadDesc (&squads)[numSquads], F f, int warpIdxStart = 0) +{ + static_assert(numSquads > 0); + if (numSquads == 1) + { + // Leaf + SquadDesc squad = squads[0]; + + if (static_cast(warpIdxStart) <= sr.warpIdx + && sr.warpIdx < static_cast(warpIdxStart + squad.warpCount())) + { + f(Squad(squad, sr)); + } + } + else + { + constexpr int mid = numSquads / 2; + // Left + int warpIdxStartMid = warpIdxStart; + for (int gi = 0; gi < mid; ++gi) + { + warpIdxStartMid += squads[gi].warpCount(); + } + if (sr.warpIdx < static_cast(warpIdxStartMid)) + { + if constexpr (0 < mid) + { + SquadDesc squadsLeft[mid]; + for (int gi = 0; gi < mid; ++gi) + { + squadsLeft[gi] = squads[gi]; + } + squadDispatch(sr, squadsLeft, f, warpIdxStart); + } + } + else + { + SquadDesc squadsRight[numSquads - mid]{}; + for (int gi = 0; gi < numSquads - mid; ++gi) + { + squadsRight[gi] = squads[mid + gi]; + } + squadDispatch(sr, squadsRight, f, warpIdxStartMid); + } + } +} + +template <::cuda::std::size_t numSquads, typename F> +_CCCL_DEVICE_API _CCCL_FORCEINLINE void +squadDispatch(SpecialRegisters sr, ::cuda::std::array squads, F f, int warpIdxStart = 0) +{ + squadDispatch(sr, squads.__elems_, f, warpIdxStart); +} +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/squad_desc.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/squad_desc.cuh new file mode 100644 index 00000000..faa787f4 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/squad/squad_desc.cuh @@ -0,0 +1,77 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +// SquadDesc - squad descriptor +// +// A squad is a collection of warps that work together in a warp-specialized +// kernel. A warp-specialized kernel has multiple squads that perform part of +// the computation. +// +// SquadDesc is a host+device constexpr-compatible class that allows describing +// the warp-specialized layout of a kernel. +// +// SquadDesc is constexpr-compatible and can be created on host and device. +struct SquadDesc +{ + int mSquadIdx = -1; + int mWarpCount = -1; + + _CCCL_HIDE_FROM_ABI constexpr SquadDesc() = default; + _CCCL_HOST_DEVICE_API constexpr SquadDesc(int squadIdx, int warpCount) noexcept + : mSquadIdx(squadIdx) + , mWarpCount(warpCount) + {} + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int warpCount() const noexcept + { + return mWarpCount; + } + + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int threadCount() const noexcept + { + return 32 * warpCount(); + } + + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator==(SquadDesc lhs, SquadDesc rhs) noexcept + { + return lhs.mSquadIdx == rhs.mSquadIdx; + } + + [[nodiscard]] _CCCL_HOST_DEVICE_API friend constexpr bool operator!=(SquadDesc lhs, SquadDesc rhs) noexcept + { + return lhs.mSquadIdx != rhs.mSquadIdx; + } +}; +// squadCountThreads +// +// Utility function to count the number of threads in an array of squad +// descriptors. It is used to launch a kernel with the correct number of +// threads. +template +[[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int squadCountThreads(const SquadDesc (&squads)[numSquads]) noexcept +{ + int sumThreads = 0; + for (int gi = 0; gi < numSquads; ++gi) + { + sumThreads += squads[gi].threadCount(); + } + return sumThreads; +} +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/sync_handler.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/sync_handler.cuh new file mode 100644 index 00000000..5b97e8de --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/sync_handler.cuh @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include + +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +// SkipSync is a tag type that is used to indicate that a SyncHandler.blockInit +// should forgo syncing. +struct SkipSync +{}; + +struct SyncHandler +{ + // reducing these values to the actually used number of resources and phases does not improve performance + static constexpr int mMaxNumResources = 10; + static constexpr int mMaxNumPhases = 4; + + // Whether barriers have been initialized. + bool mHasInitialized = false; + + // Arrays of barrier locations, number of stages, number of owning threads. + int mNextResourceHandle = 0; + int mNumStages[mMaxNumResources]{}; + int mNumPhases[mMaxNumResources]{}; + int mNumOwningThreads[mMaxNumResources][mMaxNumPhases]{}; + ::cuda::std::uint64_t* mPtrBar[mMaxNumResources][mMaxNumPhases]{}; + + constexpr SyncHandler() = default; + + // we need constant destruction for the host side single stage SMEM amount, which is only possible in C++20 +#if _CCCL_STD_VER >= 2020 + _CCCL_HOST_DEVICE_API constexpr ~SyncHandler() + { + _WS_CONSTANT_ASSERT(mHasInitialized, "SyncHandler must have been initialized at end of kernel."); + } +#endif // _CCCL_STD_VER >= 2020 + + // SyncHandler is a non-copyable, non-movable type. It must be passed by + // (mutable) reference to be useful. + SyncHandler(const SyncHandler&) = delete; // Delete copy constructor + SyncHandler(SyncHandler&&) = delete; // Delete move constructor + SyncHandler& operator=(const SyncHandler&) = delete; // Delete copy assignment + SyncHandler& operator=(const SyncHandler&&) = delete; // Delete move assignment + + // registerResource and registerPhase can be called on host and device. + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr int registerResource(int numStages) + { + _WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot register resource after SyncHandler has been initialized."); + // Avoid exceeding the max number of stages + _WS_CONSTANT_ASSERT(mNextResourceHandle < mMaxNumResources, "Cannot register more than 10 resources."); + + // Get a handle + int handle = mNextResourceHandle; + mNextResourceHandle++; + // Set the number of stages + mNumStages[handle] = numStages; + + return handle; + } + + _CCCL_HOST_DEVICE_API void constexpr registerPhase(int resourceHandle, int numOwningThreads, uint64_t* ptrBar) + { + _WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot register phase after SyncHandler has been initialized."); + _WS_CONSTANT_ASSERT(resourceHandle < mNextResourceHandle, "Invalid resource handle."); + + // Get phase index: + int curPhase = mNumPhases[resourceHandle]; + _WS_CONSTANT_ASSERT(curPhase < mMaxNumPhases, "Cannot register more phases than maximum."); + + mNumOwningThreads[resourceHandle][curPhase] = numOwningThreads; + mPtrBar[resourceHandle][curPhase] = ptrBar; + + mNumPhases[resourceHandle]++; + } + + // clusterInitSync can only be called on device. + template + _CCCL_DEVICE_API void clusterInitSync(SpecialRegisters sr, SkipSync) + { + _WS_CONSTANT_ASSERT(!mHasInitialized, "Cannot initialize SyncHandler twice."); + mHasInitialized = true; + + // All warps iterate through all resources and phases. Since all array indices have to be statically resolved by the + // SROA optimization to avoid spilling to local memory, we cannot split the iteration among warps etc. + for (int ri = 0; ri < mMaxNumResources; ri++) + { + if (ri >= mNextResourceHandle) + { + break; + } + const int resNumPhases = mNumPhases[ri]; + const int numStages = mNumStages[ri]; + + for (int pi = 0; pi < mMaxNumPhases; pi++) + { + if (pi >= resNumPhases) + { + break; + } + + uint64_t* ptrBar = mPtrBar[ri][pi]; + int numOwningThreads = mNumOwningThreads[ri][pi]; + // use block strided iteration to vectorize setup of barriers + for (int si = static_cast(sr.threadIdxX); si < numStages; si += NumThreads) + { + ::cuda::ptx::mbarrier_init(&ptrBar[si], numOwningThreads); + } + } + } + } + + template + _CCCL_DEVICE_API void clusterInitSync(SpecialRegisters sr) + { + NV_IF_TARGET(NV_PROVIDES_SM_90, ({ + clusterInitSync(sr, SkipSync{}); + __cluster_barrier_arrive_relaxed(); + __cluster_barrier_wait(); + })) + } +}; +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/values.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/values.cuh new file mode 100644 index 00000000..473639c9 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/values.cuh @@ -0,0 +1,33 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +CUB_NAMESPACE_BEGIN + +namespace detail::warpspeed +{ +enum class Stages : int +{ +}; +enum class Elems : int +{ +}; +enum class Warps : int +{ +}; +enum class Align : int +{ +}; +} // namespace detail::warpspeed + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/warpspeed.cuh b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/warpspeed.cuh new file mode 100644 index 00000000..669b3d4d --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/detail/warpspeed/warpspeed.cuh @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_adjacent_difference.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_adjacent_difference.cuh new file mode 100644 index 00000000..a762fa81 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_adjacent_difference.cuh @@ -0,0 +1,920 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceAdjacentDifference provides device-wide, parallel operations for +//! computing the differences of adjacent elements residing within +//! device-accessible memory. +//! +//! Overview +//! ++++++++++++++++++++++++++ +//! +//! - DeviceAdjacentDifference calculates the differences of adjacent elements in +//! d_input. Because the binary operation could be noncommutative, there +//! are two sets of methods. Methods named SubtractLeft subtract left element +//! ``*(i - 1)`` of input sequence from current element ``*i``. +//! Methods named ``SubtractRight`` subtract current element ``*i`` from the +//! right one ``*(i + 1)``: +//! +//! .. code-block:: c++ +//! +//! int *d_values; // [1, 2, 3, 4] +//! //... +//! int *d_subtract_left_result <-- [ 1, 1, 1, 1 ] +//! int *d_subtract_right_result <-- [ -1, -1, -1, 4 ] +//! +//! - For SubtractLeft, if the left element is out of bounds, the iterator is +//! assigned to ``*(result + (i - first))`` without modification. +//! - For SubtractRight, if the right element is out of bounds, the iterator is +//! assigned to ``*(result + (i - first))`` without modification. +//! +//! Snippet +//! ++++++++++++++++++++++++++ +//! +//! The code snippet below illustrates how to use ``DeviceAdjacentDifference`` to +//! compute the left difference between adjacent elements. +//! +//! .. code-block:: c++ +//! +//! #include +//! // or equivalently +//! +//! // Declare, allocate, and initialize device-accessible pointers +//! int num_items; // e.g., 8 +//! int *d_values; // e.g., [1, 2, 1, 2, 1, 2, 1, 2] +//! //... +//! +//! // Determine temporary device storage requirements +//! void *d_temp_storage = nullptr; +//! size_t temp_storage_bytes = 0; +//! +//! cub::DeviceAdjacentDifference::SubtractLeft( +//! d_temp_storage, temp_storage_bytes, d_values, num_items); +//! +//! // Allocate temporary storage +//! cudaMalloc(&d_temp_storage, temp_storage_bytes); +//! +//! // Run operation +//! cub::DeviceAdjacentDifference::SubtractLeft( +//! d_temp_storage, temp_storage_bytes, d_values, num_items); +//! +//! // d_values <-- [1, 1, -1, 1, -1, 1, -1, 1] +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceAdjacentDifference that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns an :cpp:struct:`cub::AdjacentDifferencePolicy`, as shown +//! in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_adjacent_difference_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin subtract-left-copy-policy-selector +//! :end-before: example-end subtract-left-copy-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_adjacent_difference_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin subtract-left-copy-tuning +//! :end-before: example-end subtract-left-copy-tuning +//! +//! @endrst +struct DeviceAdjacentDifference +{ + //! @rst + //! Subtracts the left element of each adjacent pair of elements residing within device-accessible memory + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! - Calculates the differences of adjacent elements in ``d_input``. + //! That is, ``*d_input`` is assigned to ``*d_output``, and, for each iterator ``i`` in the + //! range ``[d_input + 1, d_input + num_items)``, the result of + //! ``difference_op(*i, *(i - 1))`` is assigned to ``*(d_output + (i - d_input))``. + //! - Note that the behavior is undefined if the input and output ranges + //! overlap in any way. + //! + //! Snippet + //! ++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how to use ``DeviceAdjacentDifference`` + //! to compute the difference between adjacent elements. + //! + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! int num_items; // e.g., 8 + //! int *d_input; // e.g., [1, 2, 1, 2, 1, 2, 1, 2] + //! int *d_output; + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! + //! cub::DeviceAdjacentDifference::SubtractLeftCopy( + //! d_temp_storage, temp_storage_bytes, + //! d_input, d_output, + //! num_items, CustomDifference()); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run operation + //! cub::DeviceAdjacentDifference::SubtractLeftCopy( + //! d_temp_storage, temp_storage_bytes, + //! d_input, d_output, + //! num_items, CustomDifference()); + //! + //! // d_input <-- [1, 2, 1, 2, 1, 2, 1, 2] + //! // d_output <-- [1, 1, -1, 1, -1, 1, -1, 1] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input elements @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output elements @iterator + //! + //! @tparam DifferenceOpT + //! Its `result_type` is convertible to a type in `OutputIteratorT`'s set of `value_types`. + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @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_input + //! Beginning of the input sequence + //! + //! @param[out] d_output + //! Beginning of the output sequence + //! + //! @param[in] num_items + //! Number of items in the input sequence + //! + //! @param[in] difference_op + //! The binary function used to compute differences + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>> + static CUB_RUNTIME_FUNCTION cudaError_t SubtractLeftCopy( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_input, + OutputIteratorT d_output, + NumItemsT num_items, + DifferenceOpT difference_op = {}, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceAdjacentDifference::SubtractLeftCopy"); + + return detail::dispatch_with_env( + d_temp_storage, temp_storage_bytes, env, [&](auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::adjacent_difference::dispatch( + storage, bytes, d_input, d_output, num_items, difference_op, stream, tuning_env); + }); + } + + //! @rst + //! Subtracts the left element of each adjacent pair of elements residing within device-accessible memory. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Calculates the differences of adjacent elements in ``d_input``. That is, for + //! each iterator ``i`` in the range ``[d_input + 1, d_input + num_items)``, the + //! result of ``difference_op(*i, *(i - 1))`` is assigned to + //! ``*(d_input + (i - d_input))``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how to use ``DeviceAdjacentDifference`` + //! to compute the difference between adjacent elements. + //! + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! int num_items; // e.g., 8 + //! int *d_data; // e.g., [1, 2, 1, 2, 1, 2, 1, 2] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceAdjacentDifference::SubtractLeft( + //! d_temp_storage, temp_storage_bytes, + //! d_data, num_items, CustomDifference()); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run operation + //! cub::DeviceAdjacentDifference::SubtractLeft( + //! d_temp_storage, temp_storage_bytes, + //! d_data, num_items, CustomDifference()); + //! + //! // d_data <-- [1, 1, -1, 1, -1, 1, -1, 1] + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! **[inferred]** Random-access iterator type for reading and writing elements @iterator + //! + //! @tparam DifferenceOpT + //! Its `result_type` is convertible to a type in `RandomAccessIteratorT`'s + //! set of `value_types`. + //! + //! @tparam NumItemsT + //! **[inferred]** Type of `num_items` + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_input + //! Beginning of the input sequence and the result + //! + //! @param[in] num_items + //! Number of items in the input sequence + //! + //! @param[in] difference_op + //! The binary function used to compute differences + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>> + static CUB_RUNTIME_FUNCTION cudaError_t SubtractLeft( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorT d_input, + NumItemsT num_items, + DifferenceOpT difference_op = {}, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceAdjacentDifference::SubtractLeft"); + + return detail::dispatch_with_env( + d_temp_storage, temp_storage_bytes, env, [&](auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::adjacent_difference::dispatch( + storage, bytes, d_input, d_input, num_items, difference_op, stream, tuning_env); + }); + } + + //! @rst + //! Subtracts the right element of each adjacent pair of elements residing within device-accessible memory. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! - Calculates the right differences of adjacent elements in ``d_input``. + //! That is, ``*(d_input + num_items - 1)`` is assigned to + //! ``*(d_output + num_items - 1)``, and, for each iterator ``i`` in the range + //! ``[d_input, d_input + num_items - 1)``, the result of + //! ``difference_op(*i, *(i + 1))`` is assigned to + //! ``*(d_output + (i - d_input))``. + //! - Note that the behavior is undefined if the input and output ranges + //! overlap in any way. + //! + //! Snippet + //! ++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how to use ``DeviceAdjacentDifference`` + //! to compute the difference between adjacent elements. + //! + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! struct CustomDifference + //! { + //! template + //! __host__ DataType operator()(DataType &lhs, DataType &rhs) + //! { + //! return lhs - rhs; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! int num_items; // e.g., 8 + //! int *d_input; // e.g., [1, 2, 1, 2, 1, 2, 1, 2] + //! int *d_output; + //! .. + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceAdjacentDifference::SubtractRightCopy( + //! d_temp_storage, temp_storage_bytes, + //! d_input, d_output, num_items, CustomDifference()); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run operation + //! cub::DeviceAdjacentDifference::SubtractRightCopy( + //! d_temp_storage, temp_storage_bytes, + //! d_input, d_output, num_items, CustomDifference()); + //! + //! // d_input <-- [1, 2, 1, 2, 1, 2, 1, 2] + //! // d_data <-- [-1, 1, -1, 1, -1, 1, -1, 2] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input elements @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output elements @iterator + //! + //! @tparam DifferenceOpT + //! Its `result_type` is convertible to a type in `OutputIteratorT`'s + //! set of `value_types`. + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @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_input + //! Beginning of the input sequence + //! + //! @param[out] d_output + //! Beginning of the output sequence + //! + //! @param[in] num_items + //! Number of items in the input sequence + //! + //! @param[in] difference_op + //! The binary function used to compute differences. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>> + static CUB_RUNTIME_FUNCTION cudaError_t SubtractRightCopy( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_input, + OutputIteratorT d_output, + NumItemsT num_items, + DifferenceOpT difference_op = {}, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceAdjacentDifference::SubtractRightCopy"); + + return detail::dispatch_with_env( + d_temp_storage, temp_storage_bytes, env, [&](auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::adjacent_difference::dispatch( + storage, bytes, d_input, d_output, num_items, difference_op, stream, tuning_env); + }); + } + + //! @rst + //! Subtracts the right element of each adjacent pair of elements residing within device-accessible memory. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! Overview + //! ++++++++++++++++++++++++++ + //! + //! Calculates the right differences of adjacent elements in ``d_input``. + //! That is, for each iterator ``i`` in the range + //! ``[d_input, d_input + num_items - 1)``, the result of + //! ``difference_op(*i, *(i + 1))`` is assigned to ``*(d_input + (i - d_input))``. + //! + //! Snippet + //! ++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how to use ``DeviceAdjacentDifference`` + //! to compute the difference between adjacent elements. + //! + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! int num_items; // e.g., 8 + //! int *d_data; // e.g., [1, 2, 1, 2, 1, 2, 1, 2] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceAdjacentDifference::SubtractRight( + //! d_temp_storage, temp_storage_bytes, d_data, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run operation + //! cub::DeviceAdjacentDifference::SubtractRight( + //! d_temp_storage, temp_storage_bytes, d_data, num_items); + //! + //! // d_data <-- [-1, 1, -1, 1, -1, 1, -1, 2] + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! **[inferred]** Random-access iterator type for reading and writing elements @iterator + //! + //! @tparam DifferenceOpT + //! Its `result_type` is convertible to a type in `RandomAccessIteratorT`'s + //! set of `value_types`. + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_input + //! Beginning of the input sequence + //! + //! @param[in] num_items + //! Number of items in the input sequence + //! + //! @param[in] difference_op + //! The binary function used to compute differences + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>> + static CUB_RUNTIME_FUNCTION cudaError_t SubtractRight( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorT d_input, + NumItemsT num_items, + DifferenceOpT difference_op = {}, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceAdjacentDifference::SubtractRight"); + + return detail::dispatch_with_env( + d_temp_storage, temp_storage_bytes, env, [&](auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::adjacent_difference::dispatch( + storage, bytes, d_input, d_input, num_items, difference_op, stream, tuning_env); + }); + } + + //! @rst + //! Subtracts the left element of each adjacent pair of elements residing within device-accessible memory. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! - Calculates the differences of adjacent elements in ``d_input``. + //! That is, ``*d_input`` is assigned to ``*d_output``, and, for each iterator ``i`` in the + //! range ``[d_input + 1, d_input + num_items)``, the result of + //! ``difference_op(*i, *(i - 1))`` is assigned to ``*(d_output + (i - d_input))``. + //! - Note that the behavior is undefined if the input and output ranges + //! overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how to use ``SubtractLeftCopy`` with a custom stream + //! via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_adjacent_difference_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin subtract-left-copy-env-stream + //! :end-before: example-end subtract-left-copy-env-stream + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input elements @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output elements @iterator + //! + //! @tparam DifferenceOpT + //! **[inferred]** Binary function object type used to compute differences + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] d_input + //! Beginning of the input sequence + //! + //! @param[out] d_output + //! Beginning of the output sequence + //! + //! @param[in] num_items + //! Number of items in the input sequence + //! + //! @param[in] difference_op + //! The binary function used to compute differences + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template < + typename InputIteratorT, + typename OutputIteratorT, + typename DifferenceOpT = ::cuda::std::minus<>, + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::__indirectly_binary_invocable, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SubtractLeftCopy( + InputIteratorT d_input, + OutputIteratorT d_output, + NumItemsT num_items, + DifferenceOpT difference_op = {}, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceAdjacentDifference::SubtractLeftCopy"); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::adjacent_difference::dispatch( + storage, bytes, d_input, d_output, num_items, difference_op, stream, tuning_env); + }); + } + + //! @rst + //! Subtracts the left element of each adjacent pair of elements in-place. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Calculates the differences of adjacent elements in ``d_input``. That is, for + //! each iterator ``i`` in the range ``[d_input + 1, d_input + num_items)``, the + //! result of ``difference_op(*i, *(i - 1))`` is assigned to + //! ``*(d_input + (i - d_input))``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how to use ``SubtractLeft`` with a custom stream + //! via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_adjacent_difference_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin subtract-left-env-stream + //! :end-before: example-end subtract-left-env-stream + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! **[inferred]** Random-access iterator type for reading and writing elements @iterator + //! + //! @tparam DifferenceOpT + //! **[inferred]** Binary function object type used to compute differences + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in,out] d_input + //! Beginning of the input sequence and the result + //! + //! @param[in] num_items + //! Number of items in the input sequence + //! + //! @param[in] difference_op + //! The binary function used to compute differences + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t< + ::cuda::std::__indirectly_binary_invocable, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SubtractLeft( + RandomAccessIteratorT d_input, NumItemsT num_items, DifferenceOpT difference_op = {}, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceAdjacentDifference::SubtractLeft"); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::adjacent_difference::dispatch( + storage, bytes, d_input, d_input, num_items, difference_op, stream, tuning_env); + }); + } + + //! @rst + //! Subtracts the right element of each adjacent pair of elements residing within device-accessible memory. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! - Calculates the right differences of adjacent elements in ``d_input``. + //! That is, ``*(d_input + num_items - 1)`` is assigned to + //! ``*(d_output + num_items - 1)``, and, for each iterator ``i`` in the range + //! ``[d_input, d_input + num_items - 1)``, the result of + //! ``difference_op(*i, *(i + 1))`` is assigned to + //! ``*(d_output + (i - d_input))``. + //! - Note that the behavior is undefined if the input and output ranges + //! overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how to use ``SubtractRightCopy`` with a custom stream + //! via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_adjacent_difference_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin subtract-right-copy-env-stream + //! :end-before: example-end subtract-right-copy-env-stream + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input elements @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output elements @iterator + //! + //! @tparam DifferenceOpT + //! **[inferred]** Binary function object type used to compute differences + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] d_input + //! Beginning of the input sequence + //! + //! @param[out] d_output + //! Beginning of the output sequence + //! + //! @param[in] num_items + //! Number of items in the input sequence + //! + //! @param[in] difference_op + //! The binary function used to compute differences + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template < + typename InputIteratorT, + typename OutputIteratorT, + typename DifferenceOpT = ::cuda::std::minus<>, + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::__indirectly_binary_invocable, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SubtractRightCopy( + InputIteratorT d_input, + OutputIteratorT d_output, + NumItemsT num_items, + DifferenceOpT difference_op = {}, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceAdjacentDifference::SubtractRightCopy"); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::adjacent_difference::dispatch( + storage, bytes, d_input, d_output, num_items, difference_op, stream, tuning_env); + }); + } + + //! @rst + //! Subtracts the right element of each adjacent pair of elements in-place. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Calculates the right differences of adjacent elements in ``d_input``. + //! That is, for each iterator ``i`` in the range + //! ``[d_input, d_input + num_items - 1)``, the result of + //! ``difference_op(*i, *(i + 1))`` is assigned to ``*(d_input + (i - d_input))``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how to use ``SubtractRight`` with a custom stream + //! via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_adjacent_difference_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin subtract-right-env-stream + //! :end-before: example-end subtract-right-env-stream + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! **[inferred]** Random-access iterator type for reading and writing elements @iterator + //! + //! @tparam DifferenceOpT + //! **[inferred]** Binary function object type used to compute differences + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in,out] d_input + //! Beginning of the input sequence + //! + //! @param[in] num_items + //! Number of items in the input sequence + //! + //! @param[in] difference_op + //! The binary function used to compute differences + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t< + ::cuda::std::__indirectly_binary_invocable, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SubtractRight( + RandomAccessIteratorT d_input, NumItemsT num_items, DifferenceOpT difference_op = {}, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceAdjacentDifference::SubtractRight"); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::adjacent_difference::dispatch( + storage, bytes, d_input, d_input, num_items, difference_op, stream, tuning_env); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_batched_topk.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_batched_topk.cuh new file mode 100644 index 00000000..be153d60 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_batched_topk.cuh @@ -0,0 +1,1132 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! @file +//! cub::DeviceBatchedTopK provides device-wide, parallel operations for finding the K largest (or smallest) items +//! from many (small) segments of unordered data items residing within device-accessible memory. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include // topk::select::{min, max} +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +//! @cond +//! Shared implementation for all cub::DeviceBatchedTopK entry points. +//! +//! Validates the requested execution requirements and argument annotations, resolves the (optionally tuned) policy +//! selector from the environment, and forwards to the internal batched top-k dispatch. The selection direction is +//! threaded through as a compile-time `cuda::args::constant` so the kernel only emits the +//! requested (max OR min) code path. +//! +//! All current API-surface constraints are surfaced here as `static_assert`s so the diagnostic appears at the +//! `cub::DeviceBatchedTopK` call site rather than deep inside the kernel/agent instantiation. +template +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_batched_topk( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorItT d_keys_in, + KeyOutputIteratorItT d_keys_out, + ValueInputIteratorItT d_values_in, + ValueOutputIteratorItT d_values_out, + SegmentSizeParameterT segment_sizes, + KParameterT k, + NumSegmentsParameterT num_segments, + 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, + "Determinism should be used inside cuda::execution::require to have an effect."); + static_assert(!::cuda::std::execution::__queryable_with, + "Tie-break should be used inside cuda::execution::require to have an effect."); + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + + constexpr bool determinism_specified = + ::cuda::std::execution::__queryable_with; + constexpr bool tie_break_specified = + ::cuda::std::execution::__queryable_with; + + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + using requested_tie_break_t = + ::cuda::std::execution::__query_result_or_t; + using requested_order_t = + ::cuda::std::execution::__query_result_or_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 + || ::cuda::std::is_same_v; + constexpr bool is_non_deterministic_unsorted = + ::cuda::std::is_same_v + && ::cuda::std::is_same_v; + + static_assert(determinism_and_tie_break_paired, + "cub::DeviceBatchedTopK: determinism and tie_break requirements must be acknowledged together. Either " + "omit both to accept the defaults (cuda::execution::determinism::gpu_to_gpu and " + "cuda::execution::tie_break::prefer_smaller_index), or pass both explicitly inside " + "cuda::execution::require(...)."); + static_assert(!determinism_and_tie_break_paired || tie_break_compatible_with_determinism, + "cub::DeviceBatchedTopK: a tie_break of cuda::execution::tie_break::prefer_smaller_index or " + "prefer_larger_index pins the result set across GPUs and therefore requires " + "cuda::execution::determinism::gpu_to_gpu (it cannot be combined with run_to_run or not_guaranteed)."); + static_assert( + !determinism_and_tie_break_paired || !tie_break_compatible_with_determinism || is_non_deterministic_unsorted, + "cub::DeviceBatchedTopK currently only implements non-deterministic, unsorted output. Request it " + "explicitly with cuda::execution::require(cuda::execution::determinism::not_guaranteed, " + "cuda::execution::tie_break::unspecified, cuda::execution::output_ordering::unsorted)."); + + // --------------------------------------------------------------------------- + // Resolve the (optionally tuned) policy selector from the environment. + // --------------------------------------------------------------------------- + using key_t = cub::detail::it_value_t>; + using value_t = cub::detail::it_value_t>; + using default_policy_selector_t = batched_topk:: + policy_selector_from_types::highest>; + using tuning_env_t = + ::cuda::__call_result_or_t<::cuda::execution::__get_tuning_t, ::cuda::std::execution::env<>, EnvT>; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + + // --------------------------------------------------------------------------- + // Argument-annotation constraints surfaced at the call site. + // --------------------------------------------------------------------------- + static_assert(::cuda::args::__traits::is_single_value, + "cub::DeviceBatchedTopK currently requires a single (uniform) number of segments resolved on the " + "host; pass num_segments as a single-value annotation (e.g. cuda::args::constant or " + "cuda::args::immediate), not a per-segment sequence."); + static_assert( + ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, + "cub::DeviceBatchedTopK: segment_sizes must be a cuda::args annotation or a plain integral value " + "(taken as a uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap " + "per-segment sizes in cuda::args::deferred_sequence, or a single device-side value in " + "cuda::args::deferred."); + static_assert(::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, + "cub::DeviceBatchedTopK: k must be a cuda::args annotation or a plain integral value (taken as a " + "uniform immediate). A raw pointer or iterator is not interpreted as a sequence. Wrap a per-segment k " + "in cuda::args::deferred_sequence, or a single device-side value in cuda::args::deferred."); + static_assert( + ::cuda::args::__is_wrapper_v || ::cuda::std::is_integral_v, + "cub::DeviceBatchedTopK: num_segments must be a cuda::args annotation or a plain integral value. A " + "raw pointer or iterator is not accepted."); + + const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); + + // 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{}, + 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{}`` 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()``, 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`` or ``cuda::args::bounds()``), 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` 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 > + [[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( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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` 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 > + [[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( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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` 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 > + [[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( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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` 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 > + [[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( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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` 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 > + [[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( + 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` 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 > + [[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( + 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` 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 > + [[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( + 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` 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 > + [[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( + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, segment_sizes, k, num_segments, env); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_copy.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_copy.cuh new file mode 100644 index 00000000..786f5a2b --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_copy.cuh @@ -0,0 +1,501 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +//! @file +//! cub::DeviceCopy provides device-wide, parallel operations for copying data. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include + +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @brief cub::DeviceCopy provides device-wide, parallel operations for copying data. +//! +//! @rst +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The Batched algorithms in DeviceCopy that accept an environment can be tuned by passing a custom :ref:`policy +//! selector ` that returns a :cpp:struct:`cub::BatchedCopyPolicy`, as shown in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_copy_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin copy-batched-policy-selector +//! :end-before: example-end copy-batched-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_copy_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin copy-batched-tuning +//! :end-before: example-end copy-batched-tuning +//! @endrst +struct DeviceCopy +{ + //! @rst + //! Copies data from a batch of given source ranges to their corresponding destination ranges. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! .. note:: + //! + //! If any input range aliases any output range the behavior is undefined. + //! If any output range aliases another output range the behavior is undefined. + //! Input ranges can alias one another. + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates usage of DeviceCopy::Batched to perform a DeviceRunLength Decode operation. + //! + //! .. code-block:: c++ + //! + //! struct GetIteratorToRange + //! { + //! __host__ __device__ __forceinline__ auto operator()(uint32_t index) + //! { + //! return ::cuda::make_constant_iterator(d_data_in[index]); + //! } + //! int32_t *d_data_in; + //! }; + //! + //! struct GetPtrToRange + //! { + //! __host__ __device__ __forceinline__ auto operator()(uint32_t index) + //! { + //! return d_data_out + d_offsets[index]; + //! } + //! int32_t *d_data_out; + //! uint32_t *d_offsets; + //! }; + //! + //! struct GetRunLength + //! { + //! __host__ __device__ __forceinline__ uint32_t operator()(uint32_t index) + //! { + //! return d_offsets[index + 1] - d_offsets[index]; + //! } + //! uint32_t *d_offsets; + //! }; + //! + //! uint32_t num_ranges = 5; + //! int32_t *d_data_in; // e.g., [4, 2, 7, 3, 1] + //! int32_t *d_data_out; // e.g., [0, ... ] + //! uint32_t *d_offsets; // e.g., [0, 2, 5, 6, 9, 14] + //! + //! // Returns a constant iterator to the element of the i-th run + //! thrust::counting_iterator iota(0); + //! auto iterators_in = thrust::make_transform_iterator(iota, GetIteratorToRange{d_data_in}); + //! + //! // Returns the run length of the i-th run + //! auto sizes = thrust::make_transform_iterator(iota, GetRunLength{d_offsets}); + //! + //! // Returns pointers to the output range for each run + //! auto ptrs_out = thrust::make_transform_iterator(iota, GetPtrToRange{d_data_out, d_offsets}); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceCopy::Batched(d_temp_storage, temp_storage_bytes, iterators_in, ptrs_out, sizes, + //! num_ranges); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run batched copy algorithm (used to perform runlength decoding) + //! cub::DeviceCopy::Batched(d_temp_storage, temp_storage_bytes, iterators_in, ptrs_out, sizes, + //! num_ranges); + //! + //! // d_data_out <-- [4, 4, 2, 2, 2, 7, 3, 3, 3, 1, 1, 1, 1, 1] + //! + //! @endrst + //! + //! @tparam InputIt + //! **[inferred]** Device-accessible random-access input iterator type providing the iterators to the source ranges + //! + //! @tparam OutputIt + //! **[inferred]** Device-accessible random-access input iterator type providing the iterators to + //! the destination ranges + //! + //! @tparam SizeIteratorT + //! **[inferred]** Device-accessible random-access input iterator type providing the number of items to be + //! copied for each pair of ranges + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] input_it + //! Device-accessible iterator providing the iterators to the source ranges + //! + //! @param[in] output_it + //! Device-accessible iterator providing the iterators to the destination ranges + //! + //! @param[in] sizes + //! Device-accessible iterator providing the number of elements to be copied for each pair of ranges + //! + //! @param[in] num_ranges + //! The total number of range pairs + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t Batched( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIt input_it, + OutputIt output_it, + SizeIteratorT sizes, + ::cuda::std::int64_t num_ranges, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceCopy::Batched"); + + // Integer type large enough to hold any offset in [0, num_thread_blocks_launched), where a safe + // upper bound on num_thread_blocks_launched can be assumed to be given by + // IDIV_CEIL(num_ranges, 64) + using BlockOffsetT = uint32_t; + using default_policy_selector = detail::batch_memcpy::policy_selector; + + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::batch_memcpy::dispatch( + storage, bytes, input_it, output_it, sizes, num_ranges, stream, policy_selector); + }); + } + + //! @rst + //! Copies data from a batch of given source ranges to their corresponding destination ranges. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! .. note:: + //! + //! If any input range aliases any output range the behavior is undefined. + //! If any output range aliases another output range the behavior is undefined. + //! Input ranges can alias one another. + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates usage of DeviceCopy::Batched with an environment: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_copy_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin copy-batched-env + //! :end-before: example-end copy-batched-env + //! + //! @endrst + //! + //! @tparam InputIt + //! **[inferred]** Device-accessible random-access input iterator type providing the iterators to the source ranges + //! + //! @tparam OutputIt + //! **[inferred]** Device-accessible random-access input iterator type providing the iterators to + //! the destination ranges + //! + //! @tparam SizeIteratorT + //! **[inferred]** Device-accessible random-access input iterator type providing the number of items to be + //! copied for each pair of ranges + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] input_it + //! Device-accessible iterator providing the iterators to the source ranges + //! + //! @param[in] output_it + //! Device-accessible iterator providing the iterators to the destination ranges + //! + //! @param[in] sizes + //! Device-accessible iterator providing the number of elements to be copied for each pair of ranges + //! + //! @param[in] num_ranges + //! The total number of range pairs + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Batched( + InputIt input_it, OutputIt output_it, SizeIteratorT sizes, ::cuda::std::int64_t num_ranges, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceCopy::Batched"); + + // Integer type large enough to hold any offset in [0, num_thread_blocks_launched), where a safe + // upper bound on num_thread_blocks_launched can be assumed to be given by + // IDIV_CEIL(num_ranges, 64) + using BlockOffsetT = uint32_t; + using default_policy_selector = detail::batch_memcpy::policy_selector; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::batch_memcpy::dispatch( + storage, bytes, input_it, output_it, sizes, num_ranges, stream, policy_selector); + }); + } + + //! @rst + //! Copies data from a multidimensional source mdspan to a destination mdspan. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! This function performs a parallel copy operation between two mdspan objects with potentially different layouts but + //! identical extents. The copy operation handles arbitrary-dimensional arrays and automatically manages layout + //! transformations. + //! + //! Preconditions + //! +++++++++++++ + //! + //! * The source and destination mdspans must have identical extents (same ranks and sizes). + //! * The source and destination mdspans data handle must not be nullptr if the size is not 0. + //! * The underlying memory of the source and destination must not overlap. + //! * Both mdspans must point to device memory. + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates usage of DeviceCopy::Copy to copy between mdspans. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_copy_mdspan_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin copy-mdspan-example-op + //! :end-before: example-end copy-mdspan-example-op + //! + //! @endrst + //! + //! @tparam T_In + //! **[inferred]** The element type of the source mdspan + //! + //! @tparam Extents_In + //! **[inferred]** The extents type of the source mdspan + //! + //! @tparam Layout_In + //! **[inferred]** The layout type of the source mdspan + //! + //! @tparam Accessor_In + //! **[inferred]** The accessor type of the source mdspan + //! + //! @tparam T_Out + //! **[inferred]** The element type of the destination mdspan + //! + //! @tparam Extents_Out + //! **[inferred]** The extents type of the destination mdspan + //! + //! @tparam Layout_Out + //! **[inferred]** The layout type of the destination mdspan + //! + //! @tparam Accessor_Out + //! **[inferred]** The accessor type of the destination mdspan + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] mdspan_in + //! Source mdspan containing the data to be copied + //! + //! @param[in] mdspan_out + //! Destination mdspan where the data will be copied + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! + //! @returns + //! @rst + //! **cudaSuccess** on success, **cudaErrorInvalidValue** if mdspan extents don't match, or error code on failure + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Copy(void* d_temp_storage, + size_t& temp_storage_bytes, + ::cuda::std::mdspan mdspan_in, + ::cuda::std::mdspan mdspan_out, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return ::cudaSuccess; + } + + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceCopy::Copy"); + _CCCL_ASSERT(mdspan_in.extents() == mdspan_out.extents(), "mdspan extents must be equal"); + _CCCL_ASSERT((mdspan_in.data_handle() != nullptr && mdspan_out.data_handle() != nullptr) || mdspan_in.size() == 0, + "mdspan data handle must not be nullptr if the size is not 0"); + + // Check for memory overlap between input and output mdspans + if (mdspan_in.size() != 0) + { + auto in_start = mdspan_in.data_handle(); + auto in_end = in_start + mdspan_in.mapping().required_span_size(); + auto out_start = mdspan_out.data_handle(); + auto out_end = out_start + mdspan_out.mapping().required_span_size(); + // TODO(fbusato): replace with __are_ptrs_overlapping + _CCCL_ASSERT(!(in_end >= out_start && out_end >= in_start), "mdspan memory ranges must not overlap"); + } + + return detail::copy_mdspan::copy(mdspan_in, mdspan_out, env); + } + + //! @rst + //! Copies data from a multidimensional source mdspan to a destination mdspan. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! This function performs a parallel copy operation between two mdspan objects with potentially different layouts but + //! identical extents. The copy operation handles arbitrary-dimensional arrays and automatically manages layout + //! transformations. + //! + //! 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`` + //! + //! Preconditions + //! +++++++++++++ + //! + //! * The source and destination mdspans must have identical extents (same ranks and sizes). + //! * The source and destination mdspans data handle must not be nullptr if the size is not 0. + //! * The underlying memory of the source and destination must not overlap. + //! * Both mdspans must point to device memory. + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates usage of DeviceCopy::Copy with an environment: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_copy_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin copy-mdspan-env + //! :end-before: example-end copy-mdspan-env + //! + //! @endrst + //! + //! @tparam T_In + //! **[inferred]** The element type of the source mdspan + //! + //! @tparam Extents_In + //! **[inferred]** The extents type of the source mdspan + //! + //! @tparam Layout_In + //! **[inferred]** The layout type of the source mdspan + //! + //! @tparam Accessor_In + //! **[inferred]** The accessor type of the source mdspan + //! + //! @tparam T_Out + //! **[inferred]** The element type of the destination mdspan + //! + //! @tparam Extents_Out + //! **[inferred]** The extents type of the destination mdspan + //! + //! @tparam Layout_Out + //! **[inferred]** The layout type of the destination mdspan + //! + //! @tparam Accessor_Out + //! **[inferred]** The accessor type of the destination mdspan + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] mdspan_in + //! Source mdspan containing the data to be copied + //! + //! @param[out] mdspan_out + //! Destination mdspan where the data will be copied + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Copy(::cuda::std::mdspan mdspan_in, + ::cuda::std::mdspan mdspan_out, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceCopy::Copy"); + _CCCL_ASSERT(mdspan_in.extents() == mdspan_out.extents(), "mdspan extents must be equal"); + _CCCL_ASSERT((mdspan_in.data_handle() != nullptr && mdspan_out.data_handle() != nullptr) || mdspan_in.size() == 0, + "mdspan data handle must not be nullptr if the size is not 0"); + + // Check for memory overlap between input and output mdspans + if (mdspan_in.size() != 0) + { + auto in_start = mdspan_in.data_handle(); + auto in_end = in_start + mdspan_in.mapping().required_span_size(); + auto out_start = mdspan_out.data_handle(); + auto out_end = out_start + mdspan_out.mapping().required_span_size(); + // TODO(fbusato): replace with __are_ptrs_overlapping + _CCCL_ASSERT(!(in_end >= out_start && out_end >= in_start), "mdspan memory ranges must not overlap"); + } + + return detail::copy_mdspan::copy(mdspan_in, mdspan_out, env); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_find.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_find.cuh new file mode 100644 index 00000000..8bda4968 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_find.cuh @@ -0,0 +1,1232 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The FindIf algorithms that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::FindIfPolicy`, as shown in the +//! example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_find_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin find-if-policy-selector +//! :end-before: example-end find-if-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_find_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin find-if-tuning +//! :end-before: example-end find-if-tuning +//! +//! The ``LowerBoundSortedValues`` and ``UpperBoundSortedValues`` algorithms that accept an environment can be tuned by +//! passing a custom :ref:`policy selector ` that returns a +//! :cpp:struct:`cub::FindBoundSortedValuesPolicy`, as shown in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_find_bound_sorted_values_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin lower-bound-sorted-values-policy-selector +//! :end-before: example-end lower-bound-sorted-values-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_find_bound_sorted_values_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin lower-bound-sorted-values-tuning +//! :end-before: example-end lower-bound-sorted-values-tuning +//! +//! @endrst +struct DeviceFind +{ + //! @rst + //! Finds the first element in the input sequence that satisfies the given predicate. + //! + //! - The search terminates at the first element where the predicate evaluates to true. + //! - The index of the found element is written to ``d_out``. + //! - If no element satisfies the predicate, ``num_items`` is written to ``d_out``. + //! - The range ``[d_out, d_out + 1)`` shall not overlap ``[d_in, d_in + num_items)`` in any way. + //! - @devicestorage + //! + //! .. versionadded:: 3.3.0 + //! + //! Snippet + //! ========================================================================== + //! + //! The code snippet below illustrates the finding of the first element that satisfies the predicate. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin find-if-predicate + //! :end-before: example-end find-if-predicate + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin device-find-if + //! :end-before: example-end device-find-if + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing the result index @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Unary predicate functor type having member `bool operator()(const T &a)` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @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_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output location for the index of the found element + //! + //! @param[in] scan_op + //! Unary predicate functor for determining whether an element satisfies the search condition + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + template > + CUB_RUNTIME_FUNCTION static cudaError_t FindIf( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceFind::FindIf"); + + using OffsetT = detail::choose_offset_t; + + using default_policy_selector = detail::find::policy_selector_from_types>; + + return detail::dispatch_with_env_and_tuning( + d_temp_storage, + temp_storage_bytes, + env, + [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::find::dispatch( + storage, bytes, d_in, d_out, static_cast(num_items), scan_op, stream, policy_selector); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! For each ``value`` in ``[d_values, d_values + values_num_items)``, performs a binary search in the range + //! ``[d_range, d_range + range_num_items)``, using ``comp`` as the comparator to find the iterator to the + //! **first** element of said range which **is not** ordered **before** ``value``. + //! + //! - The range ``[d_range, d_range + range_num_items)`` must be sorted consistently with ``comp``. + //! + //! .. versionadded:: 3.3.0 + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the lower bound search. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin device-lower-bound + //! :end-before: example-end device-lower-bound + //! + //! @endrst + //! + //! @tparam RangeIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``ValuesIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam RangeNumItemsT + //! is an integral type representing the number of elements in the range to be searched. + //! + //! @tparam ValuesIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``RangeIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam ValuesNumItemsT + //! is a model of integral type representing the number of elements in the range of values to be searched for. + //! + //! @tparam OutputIteratorT + //! is a model of [Random Access Iterator], whose value type is assignable from ``RangeIteratorT``'s difference + //! type. + //! + //! @tparam CompareOpT + //! is a model of [Strict Weak Ordering], which forms a [Relation] with the value types of ``RangeIteratorT`` + //! and ``ValuesIteratorT``. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @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_range + //! Iterator to the beginning of the ordered range to be searched. + //! + //! @param[in] range_num_items + //! Number of elements in the ordered range to be searched. + //! + //! @param[in] d_values + //! Iterator to the beginning of the range of values to be searched for. + //! + //! @param[in] values_num_items + //! Number of elements in the range of values to be searched for. + //! + //! @param[out] d_output + //! Iterator to the beginning of the output range. + //! + //! @param[in] comp + //! Comparison function object which returns true if its first argument is ordered before the second in the + //! [Strict Weak Ordering] of the range to be searched. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + //! [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + //! [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + //! [Relation]: https://en.cppreference.com/w/cpp/concepts/relation + template > + CUB_RUNTIME_FUNCTION static cudaError_t LowerBound( + void* d_temp_storage, + size_t& temp_storage_bytes, + RangeIteratorT d_range, + RangeNumItemsT range_num_items, + ValuesIteratorT d_values, + ValuesNumItemsT values_num_items, + OutputIteratorT d_output, + CompareOpT comp, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceFind::LowerBound"); + + using RangeOffsetT = detail::choose_offset_t; + using ValuesOffsetT = detail::choose_offset_t; + + return detail::dispatch_with_env( + d_temp_storage, + temp_storage_bytes, + env, + [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + if (storage == nullptr) + { + bytes = 1; + return cudaSuccess; + } + + return DeviceTransform::__transform_internal( + ::cuda::std::make_tuple(d_values), + d_output, + static_cast(values_num_items), + ::cuda::always_true{}, + detail::find::make_binary_search_transform_op( + d_range, static_cast(range_num_items), comp), + ::cuda::stream_ref{stream}); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! For each ``value`` in ``[d_values, d_values + values_num_items)``, performs a binary search in the range + //! ``[d_range, d_range + range_num_items)``, + //! using ``comp`` as the comparator to find the iterator to the **first** element of said range which **is** + //! ordered **after** ``value``. + //! + //! - The range ``[d_range, d_range + range_num_items)`` must be sorted consistently with ``comp``. + //! + //! .. versionadded:: 3.3.0 + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the upper bound search. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin device-upper-bound + //! :end-before: example-end device-upper-bound + //! + //! @endrst + //! + //! @tparam RangeIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``ValuesIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam RangeNumItemsT + //! is an integral type representing the number of elements in the range to be searched. + //! + //! @tparam ValuesIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``RangeIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam ValuesNumItemsT + //! is a model of integral type representing the number of elements in the range of values to be searched for. + //! + //! @tparam OutputIteratorT + //! is a model of [Random Access Iterator], whose value type is assignable from ``RangeIteratorT``'s difference + //! type. + //! + //! @tparam CompareOpT + //! is a model of [Strict Weak Ordering], which forms a [Relation] with the value types of ``RangeIteratorT`` + //! and ``ValuesIteratorT``. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @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_range + //! Iterator to the beginning of the ordered range to be searched. + //! + //! @param[in] range_num_items + //! Number of elements in the ordered range to be searched. + //! + //! @param[in] d_values + //! Iterator to the beginning of the range of values to be searched for. + //! + //! @param[in] values_num_items + //! Number of elements in the range of values to be searched for. + //! + //! @param[out] d_output + //! Iterator to the beginning of the output range. + //! + //! @param[in] comp + //! Comparison function object which returns true if its first argument is ordered before the second in the + //! [Strict Weak Ordering] of the range to be searched. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + //! [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + //! [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + //! [Relation]: https://en.cppreference.com/w/cpp/concepts/relation + template > + CUB_RUNTIME_FUNCTION static cudaError_t UpperBound( + void* d_temp_storage, + size_t& temp_storage_bytes, + RangeIteratorT d_range, + RangeNumItemsT range_num_items, + ValuesIteratorT d_values, + ValuesNumItemsT values_num_items, + OutputIteratorT d_output, + CompareOpT comp, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceFind::UpperBound"); + + using RangeOffsetT = detail::choose_offset_t; + using ValuesOffsetT = detail::choose_offset_t; + + return detail::dispatch_with_env( + d_temp_storage, + temp_storage_bytes, + env, + [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + if (storage == nullptr) + { + bytes = 1; + return cudaSuccess; + } + + return DeviceTransform::__transform_internal( + ::cuda::std::make_tuple(d_values), + d_output, + static_cast(values_num_items), + ::cuda::always_true{}, + detail::find::make_binary_search_transform_op( + d_range, static_cast(range_num_items), comp), + ::cuda::stream_ref{stream}); + }); + } + //! @rst + //! Finds the first element in the input sequence that satisfies the given predicate. + //! + //! - The search terminates at the first element where the predicate evaluates to true. + //! - The index of the found element is written to ``d_out``. + //! - If no element satisfies the predicate, ``num_items`` is written to ``d_out``. + //! - The range ``[d_out, d_out + 1)`` shall not overlap ``[d_in, d_in + num_items)`` in any way. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the finding of the first element that satisfies the predicate. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin find-if-predicate + //! :end-before: example-end find-if-predicate + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin find-if-env + //! :end-before: example-end find-if-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing the result index @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Unary predicate functor type having member `bool operator()(const T &a)` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output location for the index of the found element + //! + //! @param[in] scan_op + //! Unary predicate functor for determining whether an element satisfies the search condition + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + FindIf(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op, NumItemsT num_items, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFind::FindIf"); + + using OffsetT = detail::choose_offset_t; + + using default_policy_selector = detail::find::policy_selector_from_types>; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::find::dispatch( + storage, bytes, d_in, d_out, static_cast(num_items), scan_op, stream, policy_selector); + }); + } + + //! @rst + //! For each ``value`` in ``[d_values, d_values + values_num_items)``, performs a binary search in the range + //! ``[d_range, d_range + range_num_items)``, using ``comp`` as the comparator to find the iterator to the + //! **first** element of said range which **is not** ordered **before** ``value``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The range ``[d_range, d_range + range_num_items)`` must be sorted consistently with ``comp``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the lower bound search. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin lower-bound-env + //! :end-before: example-end lower-bound-env + //! + //! @endrst + //! + //! @tparam RangeIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``ValuesIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam RangeNumItemsT + //! is an integral type representing the number of elements in the range to be searched. + //! + //! @tparam ValuesIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``RangeIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam ValuesNumItemsT + //! is a model of integral type representing the number of elements in the range of values to be searched for. + //! + //! @tparam OutputIteratorT + //! is a model of [Random Access Iterator], whose value type is assignable from ``RangeIteratorT``'s difference + //! type. + //! + //! @tparam CompareOpT + //! is a model of [Strict Weak Ordering], which forms a [Relation] with the value types of ``RangeIteratorT`` + //! and ``ValuesIteratorT``. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_range + //! Iterator to the beginning of the ordered range to be searched. + //! + //! @param[in] range_num_items + //! Number of elements in the ordered range to be searched. + //! + //! @param[in] d_values + //! Iterator to the beginning of the range of values to be searched for. + //! + //! @param[in] values_num_items + //! Number of elements in the range of values to be searched for. + //! + //! @param[out] d_output + //! Iterator to the beginning of the output range. + //! + //! @param[in] comp + //! Comparison function object which returns true if its first argument is ordered before the second in the + //! [Strict Weak Ordering] of the range to be searched. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + //! [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + //! [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + //! [Relation]: https://en.cppreference.com/w/cpp/concepts/relation + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t LowerBound( + RangeIteratorT d_range, + RangeNumItemsT range_num_items, + ValuesIteratorT d_values, + ValuesNumItemsT values_num_items, + OutputIteratorT d_output, + CompareOpT comp, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFind::LowerBound"); + + using RangeOffsetT = detail::choose_offset_t; + using ValuesOffsetT = detail::choose_offset_t; + + return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + if (storage == nullptr) + { + bytes = 1; + return cudaSuccess; + } + + return DeviceTransform::__transform_internal( + ::cuda::std::make_tuple(d_values), + d_output, + static_cast(values_num_items), + ::cuda::always_true{}, + detail::find::make_binary_search_transform_op( + d_range, static_cast(range_num_items), comp), + ::cuda::stream_ref{stream}); + }); + } + + //! @rst + //! For each ``value`` in ``[d_values, d_values + values_num_items)``, performs a binary search in the range + //! ``[d_range, d_range + range_num_items)``, + //! using ``comp`` as the comparator to find the iterator to the **first** element of said range which **is** + //! ordered **after** ``value``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The range ``[d_range, d_range + range_num_items)`` must be sorted consistently with ``comp``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the upper bound search. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin upper-bound-env + //! :end-before: example-end upper-bound-env + //! + //! @endrst + //! + //! @tparam RangeIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``ValuesIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam RangeNumItemsT + //! is an integral type representing the number of elements in the range to be searched. + //! + //! @tparam ValuesIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``RangeIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam ValuesNumItemsT + //! is a model of integral type representing the number of elements in the range of values to be searched for. + //! + //! @tparam OutputIteratorT + //! is a model of [Random Access Iterator], whose value type is assignable from ``RangeIteratorT``'s difference + //! type. + //! + //! @tparam CompareOpT + //! is a model of [Strict Weak Ordering], which forms a [Relation] with the value types of ``RangeIteratorT`` + //! and ``ValuesIteratorT``. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_range + //! Iterator to the beginning of the ordered range to be searched. + //! + //! @param[in] range_num_items + //! Number of elements in the ordered range to be searched. + //! + //! @param[in] d_values + //! Iterator to the beginning of the range of values to be searched for. + //! + //! @param[in] values_num_items + //! Number of elements in the range of values to be searched for. + //! + //! @param[out] d_output + //! Iterator to the beginning of the output range. + //! + //! @param[in] comp + //! Comparison function object which returns true if its first argument is ordered before the second in the + //! [Strict Weak Ordering] of the range to be searched. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + //! [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + //! [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + //! [Relation]: https://en.cppreference.com/w/cpp/concepts/relation + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t UpperBound( + RangeIteratorT d_range, + RangeNumItemsT range_num_items, + ValuesIteratorT d_values, + ValuesNumItemsT values_num_items, + OutputIteratorT d_output, + CompareOpT comp, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFind::UpperBound"); + + using RangeOffsetT = detail::choose_offset_t; + using ValuesOffsetT = detail::choose_offset_t; + + return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + if (storage == nullptr) + { + bytes = 1; + return cudaSuccess; + } + + return DeviceTransform::__transform_internal( + ::cuda::std::make_tuple(d_values), + d_output, + static_cast(values_num_items), + ::cuda::always_true{}, + detail::find::make_binary_search_transform_op( + d_range, static_cast(range_num_items), comp), + ::cuda::stream_ref{stream}); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Accelerated variant of :cpp:func:`LowerBound` that exploits the additional + //! precondition that ``[d_values, d_values + values_num_items)`` is also + //! sorted consistently with ``comp``. + //! + //! For each ``value`` in ``[d_values, d_values + values_num_items)``, + //! performs a search in ``[d_range, d_range + range_num_items)`` to find the + //! iterator to the first element that is **not ordered before** ``value``. + //! + //! Because both sequences are sorted, the algorithm uses the Merge-Path + //! algorithm (Oded et al., IPDPS 2012) to partition the combined traversal + //! across thread blocks, achieving O(N+M) total device work rather than the + //! O(M log N) of independent binary searches. + //! + //! - Both ``[d_range, d_range + range_num_items)`` **and** + //! ``[d_values, d_values + values_num_items)`` must be sorted consistently + //! with ``comp``. + //! - @devicestorage + //! + //! .. versionadded:: 3.5.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! @endrst + //! + //! @tparam RangeIteratorT + //! is a model of [Random Access Iterator], whose value type forms a + //! [Relation] with the value type of ``ValuesIteratorT`` via ``CompareOpT``. + //! + //! @tparam RangeNumItemsT + //! is an integral type representing the number of elements in the range. + //! + //! @tparam ValuesIteratorT + //! is a model of [Random Access Iterator], whose value type forms a + //! [Relation] with the value type of ``RangeIteratorT`` via ``CompareOpT``. + //! + //! @tparam ValuesNumItemsT + //! is an integral type representing the number of values to search for. + //! + //! @tparam OutputIteratorT + //! is a model of [Random Access Iterator] whose value type is assignable + //! from ``RangeIteratorT``'s difference type. + //! + //! @tparam CompareOpT + //! is a model of [Strict Weak Ordering] over the value types of both + //! iterator types. + //! + //! @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_range + //! Iterator to the beginning of the ordered haystack range. + //! + //! @param[in] range_num_items + //! Number of elements in the haystack range. + //! + //! @param[in] d_values + //! Iterator to the beginning of the sorted range of needles. + //! + //! @param[in] values_num_items + //! Number of needle elements. + //! + //! @param[out] d_output + //! Iterator to the beginning of the output range. + //! + //! @param[in] comp + //! Comparison function object (Strict Weak Ordering). + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + //! + //! [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + //! [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + //! [Relation]: https://en.cppreference.com/w/cpp/concepts/relation + template + CUB_RUNTIME_FUNCTION static cudaError_t LowerBoundSortedValues( + void* d_temp_storage, + size_t& temp_storage_bytes, + RangeIteratorT d_range, + RangeNumItemsT range_num_items, + ValuesIteratorT d_values, + ValuesNumItemsT values_num_items, + OutputIteratorT d_output, + CompareOpT comp, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceFind::LowerBoundSortedValues"); + + using RangeOffsetT = detail::choose_offset_t; + using ValuesOffsetT = detail::choose_offset_t; + using OffsetT = ::cuda::std::common_type_t; + + return detail::find_bound_sorted_values::dispatch( + d_temp_storage, + temp_storage_bytes, + d_range, + static_cast(range_num_items), + d_values, + static_cast(values_num_items), + d_output, + comp, + stream); + } + + //! @rst + //! Accelerated variant of :cpp:func:`LowerBound` that exploits the additional + //! precondition that ``[d_values, d_values + values_num_items)`` is also + //! sorted consistently with ``comp``. + //! + //! .. versionadded:: 3.5.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Both ``[d_range, d_range + range_num_items)`` **and** + //! ``[d_values, d_values + values_num_items)`` must be sorted consistently + //! with ``comp``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the lower bound search. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_bound_sorted_values_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin lower-bound-sorted-values-env + //! :end-before: example-end lower-bound-sorted-values-env + //! + //! @endrst + //! + //! @tparam RangeIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``ValuesIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam RangeNumItemsT + //! is an integral type representing the number of elements in the range to be searched. + //! + //! @tparam ValuesIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``RangeIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam ValuesNumItemsT + //! is a model of integral type representing the number of elements in the range of values to be searched for. + //! + //! @tparam OutputIteratorT + //! is a model of [Random Access Iterator], whose value type is assignable from ``RangeIteratorT``'s difference + //! type. + //! + //! @tparam CompareOpT + //! is a model of [Strict Weak Ordering], which forms a [Relation] with the value types of ``RangeIteratorT`` + //! and ``ValuesIteratorT``. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_range + //! Iterator to the beginning of the ordered haystack range. + //! + //! @param[in] range_num_items + //! Number of elements in the haystack range. + //! + //! @param[in] d_values + //! Iterator to the beginning of the sorted range of needles. + //! + //! @param[in] values_num_items + //! Number of needle elements. + //! + //! @param[out] d_output + //! Iterator to the beginning of the output range. + //! + //! @param[in] comp + //! Comparison function object (Strict Weak Ordering). + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + //! [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + //! [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + //! [Relation]: https://en.cppreference.com/w/cpp/concepts/relation + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t LowerBoundSortedValues( + RangeIteratorT d_range, + RangeNumItemsT range_num_items, + ValuesIteratorT d_values, + ValuesNumItemsT values_num_items, + OutputIteratorT d_output, + CompareOpT comp, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFind::LowerBoundSortedValues"); + + using RangeOffsetT = detail::choose_offset_t; + using ValuesOffsetT = detail::choose_offset_t; + using OffsetT = ::cuda::std::common_type_t; + + using default_policy_selector = + detail::find_bound_sorted_values::policy_selector_from_types, + detail::it_value_t>; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::find_bound_sorted_values::dispatch( + storage, + bytes, + d_range, + static_cast(range_num_items), + d_values, + static_cast(values_num_items), + d_output, + comp, + stream, + policy_selector); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Accelerated variant of :cpp:func:`UpperBound` that exploits the additional + //! precondition that ``[d_values, d_values + values_num_items)`` is also + //! sorted consistently with ``comp``. + //! + //! For each ``value`` in ``[d_values, d_values + values_num_items)``, + //! performs a search in ``[d_range, d_range + range_num_items)`` to find the + //! iterator to the first element that is **ordered after** ``value``. + //! + //! Because both sequences are sorted, the algorithm uses the Merge-Path + //! algorithm (Oded et al., IPDPS 2012) to partition the combined traversal + //! across thread blocks, achieving O(N+M) total device work rather than the + //! O(M log N) of independent binary searches. + //! + //! - Both ``[d_range, d_range + range_num_items)`` **and** + //! ``[d_values, d_values + values_num_items)`` must be sorted consistently + //! with ``comp``. + //! - @devicestorage + //! + //! .. versionadded:: 3.5.0 + //! + //! @endrst + //! + //! @tparam RangeIteratorT + //! is a model of [Random Access Iterator], whose value type forms a + //! [Relation] with the value type of ``ValuesIteratorT`` via ``CompareOpT``. + //! + //! @tparam RangeNumItemsT + //! is an integral type representing the number of elements in the range. + //! + //! @tparam ValuesIteratorT + //! is a model of [Random Access Iterator], whose value type forms a + //! [Relation] with the value type of ``RangeIteratorT`` via ``CompareOpT``. + //! + //! @tparam ValuesNumItemsT + //! is an integral type representing the number of values to search for. + //! + //! @tparam OutputIteratorT + //! is a model of [Random Access Iterator] whose value type is assignable + //! from ``RangeIteratorT``'s difference type. + //! + //! @tparam CompareOpT + //! is a model of [Strict Weak Ordering] over the value types of both + //! iterator types. + //! + //! @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_range + //! Iterator to the beginning of the ordered haystack range. + //! + //! @param[in] range_num_items + //! Number of elements in the haystack range. + //! + //! @param[in] d_values + //! Iterator to the beginning of the sorted range of needles. + //! + //! @param[in] values_num_items + //! Number of needle elements. + //! + //! @param[out] d_output + //! Iterator to the beginning of the output range. + //! + //! @param[in] comp + //! Comparison function object (Strict Weak Ordering). + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + //! + //! [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + //! [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + //! [Relation]: https://en.cppreference.com/w/cpp/concepts/relation + template + CUB_RUNTIME_FUNCTION static cudaError_t UpperBoundSortedValues( + void* d_temp_storage, + size_t& temp_storage_bytes, + RangeIteratorT d_range, + RangeNumItemsT range_num_items, + ValuesIteratorT d_values, + ValuesNumItemsT values_num_items, + OutputIteratorT d_output, + CompareOpT comp, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceFind::UpperBoundSortedValues"); + + using RangeOffsetT = detail::choose_offset_t; + using ValuesOffsetT = detail::choose_offset_t; + using OffsetT = ::cuda::std::common_type_t; + + return detail::find_bound_sorted_values::dispatch( + d_temp_storage, + temp_storage_bytes, + d_range, + static_cast(range_num_items), + d_values, + static_cast(values_num_items), + d_output, + comp, + stream); + } + + //! @rst + //! Accelerated variant of :cpp:func:`UpperBound` that exploits the additional + //! precondition that ``[d_values, d_values + values_num_items)`` is also + //! sorted consistently with ``comp``. + //! + //! .. versionadded:: 3.5.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Both ``[d_range, d_range + range_num_items)`` **and** + //! ``[d_values, d_values + values_num_items)`` must be sorted consistently + //! with ``comp``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the upper bound search. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_find_bound_sorted_values_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin upper-bound-sorted-values-env + //! :end-before: example-end upper-bound-sorted-values-env + //! + //! @endrst + //! + //! @tparam RangeIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``ValuesIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam RangeNumItemsT + //! is an integral type representing the number of elements in the range to be searched. + //! + //! @tparam ValuesIteratorT + //! is a model of [Random Access Iterator], whose value type forms a [Relation] with the value type of + //! ``RangeIteratorT`` using ``CompareOpT`` as the predicate. + //! + //! @tparam ValuesNumItemsT + //! is a model of integral type representing the number of elements in the range of values to be searched for. + //! + //! @tparam OutputIteratorT + //! is a model of [Random Access Iterator], whose value type is assignable from ``RangeIteratorT``'s difference + //! type. + //! + //! @tparam CompareOpT + //! is a model of [Strict Weak Ordering], which forms a [Relation] with the value types of ``RangeIteratorT`` + //! and ``ValuesIteratorT``. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_range + //! Iterator to the beginning of the ordered haystack range. + //! + //! @param[in] range_num_items + //! Number of elements in the haystack range. + //! + //! @param[in] d_values + //! Iterator to the beginning of the sorted range of needles. + //! + //! @param[in] values_num_items + //! Number of needle elements. + //! + //! @param[out] d_output + //! Iterator to the beginning of the output range. + //! + //! @param[in] comp + //! Comparison function object (Strict Weak Ordering). + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + //! [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + //! [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + //! [Relation]: https://en.cppreference.com/w/cpp/concepts/relation + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t UpperBoundSortedValues( + RangeIteratorT d_range, + RangeNumItemsT range_num_items, + ValuesIteratorT d_values, + ValuesNumItemsT values_num_items, + OutputIteratorT d_output, + CompareOpT comp, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFind::UpperBoundSortedValues"); + + using RangeOffsetT = detail::choose_offset_t; + using ValuesOffsetT = detail::choose_offset_t; + using OffsetT = ::cuda::std::common_type_t; + + using default_policy_selector = + detail::find_bound_sorted_values::policy_selector_from_types, + detail::it_value_t>; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::find_bound_sorted_values::dispatch( + storage, + bytes, + d_range, + static_cast(range_num_items), + d_values, + static_cast(values_num_items), + d_output, + comp, + stream, + policy_selector); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_for.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_for.cuh new file mode 100644 index 00000000..a521724d --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_for.cuh @@ -0,0 +1,1375 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceFor provides device-wide, parallel operations for iterating over data elements. +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceFor that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::ForPolicy`, as shown in the +//! example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin bulk-policy-selector +//! :end-before: example-end bulk-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin bulk-tuning +//! :end-before: example-end bulk-tuning +//! +//! @endrst +struct DeviceFor +{ + //! `__op_wrapper_t` turns bulk into a for-each operation by wrapping the user-provided unary operator. + template + struct __op_wrapper_t + { + static_assert(::cuda::std::is_integral_v); + + RandomAccessIteratorT input; + OpT op; + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(OffsetT i) + { + // Dereferencing `thrust::device_vector` iterators returns a `thrust::device_reference` + // instead of `T`. Since user-provided operator expects `T` as an argument, we need to unwrap. + (void) op(THRUST_NS_QUALIFIER::raw_reference_cast(*(input + i))); + } + }; + + //! `__op_wrapper_vectorized_t` turns bulk into a for-each-copy operation. + //! `__op_wrapper_vectorized_t` is similar to `op_wrapper_t` but does not provide any guarantees about + //! address of the input parameter. `OpT` might be given a copy of the value or an actual reference + //! to the input iterator value (depending on the alignment of input iterator) + template + struct __op_wrapper_vectorized_t + { + static_assert(::cuda::std::is_integral_v); + + const T* input; // Raw pointer to the input data + OpT op; // User-provided operator + OffsetT partially_filled_vector_id; // Index of the vector that doesn't have all elements + OffsetT num_items; // Total number of non-vectorized items + + // TODO Can be extracted into tuning + constexpr static int vec_size = 4; + + // Type of the vector that is used to load the input data + using vector_t = typename CubVector::Type; + + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(OffsetT i) + { + // Surrounding `Bulk` call doesn't invoke this operator on invalid indices, so we don't need to + // check for out-of-bounds access here. + if (i == partially_filled_vector_id) + { // Case of partially filled vector + for (OffsetT j = i * vec_size; j < num_items; j++) + { + (void) op(input[j]); + } + } + else + { // Case of fully filled vector + const vector_t vec = *reinterpret_cast(input + vec_size * i); + + _CCCL_PRAGMA_UNROLL_FULL() + for (int j = 0; j < vec_size; j++) + { + (void) op(*(reinterpret_cast(&vec) + j)); + } + } + } + }; + + template + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t __bulk(OffsetT num_items, OpT op, const EnvT& env = {}) + { + auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env); + [[maybe_unused]] const auto tuning_env = + ::cuda::__call_or(::cuda::execution::__get_tuning, ::cuda::std::execution::env<>{}, env); + using default_policy_selector = detail::for_each::policy_selector; + using policy_selector = + ::cuda::std::execution::__query_result_or_t; + return detail::for_each::dispatch(num_items, op, stream.get(), policy_selector{}); + } + + template + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + __for_each_n(RandomAccessIteratorT first, NumItemsT num_items, OpT op, const EnvT& env) + { + // We tried to detect if we can still use vectorization from the non-Copy CUB APIs, but it's disabled for now: + constexpr bool allow_vectorization = + (AllowCopy + /*|| detail::for_each::can_regain_copy_freedom, OpT>::value*/); + + if constexpr (allow_vectorization && THRUST_NS_QUALIFIER::is_contiguous_iterator_v) + { + auto* unwrapped_first = THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(first); + using wrapped_op_t = __op_wrapper_vectorized_t>; + + if (::cuda::std::is_sufficiently_aligned(unwrapped_first)) + { // Vectorize loads + const NumItemsT num_vec_items = ::cuda::ceil_div(num_items, wrapped_op_t::vec_size); + return __bulk( + num_vec_items, wrapped_op_t{unwrapped_first, op, num_items / wrapped_op_t::vec_size, num_items}, env); + } + } + + return __bulk(num_items, __op_wrapper_t{first, op}, env); + } + + template + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t __for_each_n( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorT first, + NumItemsT num_items, + OpT op, + cudaStream_t stream = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + return __for_each_n(first, num_items, op, stream); + } + +public: + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each index in the provided shape + //! The algorithm is similar to + //! `bulk `_ + //! from P2300. + //! + //! .. versionadded:: 2.4.0 + //! First appears in CUDA Toolkit 12.5. + //! + //! - The return value of ``op``, if any, is ignored. + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use Bulk to square each element in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-square-t + //! :end-before: example-end bulk-square-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-temp-storage + //! :end-before: example-end bulk-temp-storage + //! + //! @endrst + //! + //! @tparam ShapeT + //! is an integral type + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] shape + //! Shape of the index space to iterate over + //! + //! @param[in] op + //! Function object to apply to each index in the index space + //! + //! @param[in] stream + //! CUDA stream to launch kernels within. Default stream is `0`. + template + CUB_RUNTIME_FUNCTION static cudaError_t + Bulk(void* d_temp_storage, size_t& temp_storage_bytes, ShapeT shape, OpT op, cudaStream_t stream = {}) + { + static_assert(::cuda::std::is_integral_v, "ShapeT must be an integral type"); + + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return Bulk(shape, op, stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each element in the range ``[first, first + num_items)`` + //! + //! .. versionadded:: 2.4.0 + //! First appears in CUDA Toolkit 12.5. + //! + //! - The return value of ``op``, if any, is ignored. + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachN` to square each element in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-square-ref-t + //! :end-before: example-end bulk-square-ref-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-n-temp-storage + //! :end-before: example-end for-each-n-temp-storage + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! is a model of Random Access Iterator whose value type is convertible to `op`'s argument type. + //! + //! @tparam NumItemsT + //! is an integral type representing the number of elements to iterate over + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] first + //! The beginning of the sequence + //! + //! @param[in] num_items + //! Number of elements to iterate over + //! + //! @param[in] op + //! Function object to apply to each element in the range + //! + //! @param[in] stream + //! CUDA stream to launch kernels within. Default stream is `0`. + template + CUB_RUNTIME_FUNCTION static cudaError_t ForEachN( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorT first, + NumItemsT num_items, + OpT op, + cudaStream_t stream = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return ForEachN(first, num_items, op, stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each element in the range ``[first, last)`` + //! + //! .. versionadded:: 2.4.0 + //! First appears in CUDA Toolkit 12.5. + //! + //! - The return value of ``op``, if any, is ignored. + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEach` to square each element in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-square-ref-t + //! :end-before: example-end bulk-square-ref-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-temp-storage + //! :end-before: example-end for-each-temp-storage + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! is a model of Random Access Iterator whose value type is convertible to `op`'s argument type. + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] first + //! The beginning of the sequence + //! + //! @param[in] last + //! The end of the sequence + //! + //! @param[in] op + //! Function object to apply to each element in the range + //! + //! @param[in] stream + //! CUDA stream to launch kernels within. Default stream is `0`. + template + CUB_RUNTIME_FUNCTION static cudaError_t ForEach( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorT first, + RandomAccessIteratorT last, + OpT op, + cudaStream_t stream = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return ForEach(first, last, op, stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each element in the range ``[first, first + num_items)``. + //! Unlike the ``ForEachN`` algorithm, ``ForEachCopyN`` is allowed to invoke ``op`` on copies of the elements. + //! This relaxation allows ``ForEachCopyN`` to vectorize loads. + //! + //! .. versionadded:: 2.4.0 + //! First appears in CUDA Toolkit 12.5. + //! + //! - Allowed to invoke ``op`` on copies of the elements + //! - The return value of ``op``, if any, is ignored. + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachCopyN` to count odd elements in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-odd-count-t + //! :end-before: example-end bulk-odd-count-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-copy-n-temp-storage + //! :end-before: example-end for-each-copy-n-temp-storage + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! is a model of Random Access Iterator whose value type is convertible to `op`'s argument type. + //! + //! @tparam NumItemsT + //! is an integral type representing the number of elements to iterate over + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] first + //! The beginning of the sequence + //! + //! @param[in] num_items + //! Number of elements to iterate over + //! + //! @param[in] op + //! Function object to apply to a copy of each element in the range + //! + //! @param[in] stream + //! CUDA stream to launch kernels within. Default stream is `0`. + template + CUB_RUNTIME_FUNCTION static cudaError_t ForEachCopyN( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorT first, + NumItemsT num_items, + OpT op, + cudaStream_t stream = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return ForEachCopyN(first, num_items, op, stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each element in the range ``[first, last)``. + //! Unlike the ``ForEach`` algorithm, ``ForEachCopy`` is allowed to invoke ``op`` on copies of the elements. + //! This relaxation allows ``ForEachCopy`` to vectorize loads. + //! + //! .. versionadded:: 2.4.0 + //! First appears in CUDA Toolkit 12.5. + //! + //! - Allowed to invoke ``op`` on copies of the elements + //! - The return value of ``op``, if any, is ignored. + //! - @devicestorage + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachCopy` to count odd elements in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-odd-count-t + //! :end-before: example-end bulk-odd-count-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-copy-temp-storage + //! :end-before: example-end for-each-copy-temp-storage + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! is a model of Random Access Iterator whose value type is convertible to `op`'s argument type. + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] first + //! The beginning of the sequence + //! + //! @param[in] last + //! The end of the sequence + //! + //! @param[in] op + //! Function object to apply to a copy of each element in the range + //! + //! @param[in] stream + //! CUDA stream to launch kernels within. Default stream is `0`. + template + CUB_RUNTIME_FUNCTION static cudaError_t ForEachCopy( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorT first, + RandomAccessIteratorT last, + OpT op, + cudaStream_t stream = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return ForEachCopy(first, last, op, stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each index in the provided shape + //! The algorithm is similar to + //! `bulk `_ + //! from P2300. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! This is an environment-based API that allows customization of: + //! + //! - Stream: Query via ``cuda::get_stream`` + //! + //! .. note:: + //! + //! The return value of ``op``, if any, is ignored. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use Bulk to square each element in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-square-t + //! :end-before: example-end bulk-square-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-wo-temp-storage + //! :end-before: example-end bulk-wo-temp-storage + //! + //! Environment Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use Bulk with a custom stream via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-square-env-t + //! :end-before: example-end bulk-square-env-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-env + //! :end-before: example-end bulk-env + //! + //! @endrst + //! + //! @tparam ShapeT + //! is an integral type + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] shape + //! Shape of the index space to iterate over + //! + //! @param[in] op + //! Function object to apply to each index in the index space + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t Bulk(ShapeT shape, OpT op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFor::Bulk"); + static_assert(::cuda::std::is_integral_v, "ShapeT must be an integral type"); + if (shape == 0) + { + return cudaSuccess; + } + return __bulk(shape, op, env); + } + + // we need this so the previous overload is not ambiguous with the next one + static_assert(!::cuda::std::is_convertible_v<::cuda::stream_ref, cudaStream_t>); + + // We keep this overload around to support types that are convertible to `cudaStream_t` but not copyable + template + CUB_RUNTIME_FUNCTION static cudaError_t Bulk(ShapeT shape, OpT op, cudaStream_t stream) + { + return Bulk(shape, op, ::cuda::stream_ref{stream}); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each element in the range ``[first, first + num_items)`` + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! This is an environment-based API that allows customization of: + //! + //! - Stream: Query via ``cuda::get_stream`` + //! + //! .. note:: + //! + //! The return value of ``op``, if any, is ignored. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachN` to square each element in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-square-ref-t + //! :end-before: example-end bulk-square-ref-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-n-wo-temp-storage + //! :end-before: example-end for-each-n-wo-temp-storage + //! + //! Environment Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachN` with a custom stream via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin square-ref-env-t + //! :end-before: example-end square-ref-env-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-n-env + //! :end-before: example-end for-each-n-env + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! is a model of Random Access Iterator whose value type is convertible to `op`'s argument type. + //! + //! @tparam NumItemsT + //! is an integral type representing the number of elements to iterate over + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] first + //! The beginning of the sequence + //! + //! @param[in] num_items + //! Number of elements to iterate over + //! + //! @param[in] op + //! Function object to apply to each element in the range + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachN(RandomAccessIteratorT first, NumItemsT num_items, OpT op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFor::ForEachN"); + return __for_each_n(first, num_items, op, env); + } + + // We keep this overload around to support types that are convertible to `cudaStream_t` but not copyable + template + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachN(RandomAccessIteratorT first, NumItemsT num_items, OpT op, cudaStream_t stream) + { + return ForEachN(first, num_items, op, ::cuda::stream_ref{stream}); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each element in the range ``[first, last)`` + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! This is an environment-based API that allows customization of: + //! + //! - Stream: Query via ``cuda::get_stream`` + //! + //! .. note:: + //! + //! The return value of ``op``, if any, is ignored. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEach` to square each element in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-square-ref-t + //! :end-before: example-end bulk-square-ref-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-wo-temp-storage + //! :end-before: example-end for-each-wo-temp-storage + //! + //! Environment Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEach` with a custom stream via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin square-ref-env-t + //! :end-before: example-end square-ref-env-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-env + //! :end-before: example-end for-each-env + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! is a model of Random Access Iterator whose value type is convertible to `op`'s argument type. + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] first + //! The beginning of the sequence + //! + //! @param[in] last + //! The end of the sequence + //! + //! @param[in] op + //! Function object to apply to each element in the range + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t + ForEach(RandomAccessIteratorT first, RandomAccessIteratorT last, OpT op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFor::ForEach"); + using offset_t = detail::it_difference_t; + const auto num_items = static_cast(::cuda::std::distance(first, last)); + return __for_each_n(first, num_items, op, env); + } + + // We keep this overload around to support types that are convertible to `cudaStream_t` but not copyable + template + CUB_RUNTIME_FUNCTION static cudaError_t + ForEach(RandomAccessIteratorT first, RandomAccessIteratorT last, OpT op, cudaStream_t stream) + { + return ForEach(first, last, op, ::cuda::stream_ref{stream}); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each element in the range ``[first, first + num_items)``. + //! Unlike the ``ForEachN`` algorithm, ``ForEachCopyN`` is allowed to invoke ``op`` on copies of the elements. + //! This relaxation allows ``ForEachCopyN`` to vectorize loads. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! This is an environment-based API that allows customization of: + //! + //! - Stream: Query via ``cuda::get_stream`` + //! + //! - Allowed to invoke ``op`` on copies of the elements + //! - The return value of ``op``, if any, is ignored. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachCopyN` to count odd elements in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-odd-count-t + //! :end-before: example-end bulk-odd-count-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-copy-n-wo-temp-storage + //! :end-before: example-end for-each-copy-n-wo-temp-storage + //! + //! Environment Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachCopyN` with a custom stream via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin odd-count-env-t + //! :end-before: example-end odd-count-env-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-copy-n-env + //! :end-before: example-end for-each-copy-n-env + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! is a model of Random Access Iterator whose value type is convertible to `op`'s argument type. + //! + //! @tparam NumItemsT + //! is an integral type representing the number of elements to iterate over + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] first + //! The beginning of the sequence + //! + //! @param[in] num_items + //! Number of elements to iterate over + //! + //! @param[in] op + //! Function object to apply to a copy of each element in the range + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachCopyN(RandomAccessIteratorT first, NumItemsT num_items, OpT op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFor::ForEachCopyN"); + return __for_each_n(first, num_items, op, env); + } + + // We keep this overload around to support types that are convertible to `cudaStream_t` but not copyable + template + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachCopyN(RandomAccessIteratorT first, NumItemsT num_items, OpT op, cudaStream_t stream) + { + return ForEachCopyN(first, num_items, op, ::cuda::stream_ref{stream}); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Applies the function object ``op`` to each element in the range ``[first, last)``. + //! Unlike the ``ForEach`` algorithm, ``ForEachCopy`` is allowed to invoke ``op`` on copies of the elements. + //! This relaxation allows ``ForEachCopy`` to vectorize loads. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! This is an environment-based API that allows customization of: + //! + //! - Stream: Query via ``cuda::get_stream`` + //! + //! - Allowed to invoke ``op`` on copies of the elements + //! - The return value of ``op``, if any, is ignored. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachCopy` to count odd elements in a device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin bulk-odd-count-t + //! :end-before: example-end bulk-odd-count-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-copy-wo-temp-storage + //! :end-before: example-end for-each-copy-wo-temp-storage + //! + //! Environment Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use `ForEachCopy` with a custom stream via an environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin odd-count-env-t + //! :end-before: example-end odd-count-env-t + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-copy-env + //! :end-before: example-end for-each-copy-env + //! + //! @endrst + //! + //! @tparam RandomAccessIteratorT + //! is a model of Random Access Iterator whose value type is convertible to `op`'s argument type. + //! + //! @tparam OpT + //! is a model of [Unary Function](https://en.cppreference.com/w/cpp/utility/functional/unary_function) + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] first + //! The beginning of the sequence + //! + //! @param[in] last + //! The end of the sequence + //! + //! @param[in] op + //! Function object to apply to a copy of each element in the range + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachCopy(RandomAccessIteratorT first, RandomAccessIteratorT last, OpT op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFor::ForEachCopy"); + using offset_t = detail::it_difference_t; + const auto num_items = static_cast(::cuda::std::distance(first, last)); + return __for_each_n(first, num_items, op, env); + } + + // We keep this overload around to support types that are convertible to `cudaStream_t` but not copyable + template + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachCopy(RandomAccessIteratorT first, RandomAccessIteratorT last, OpT op, cudaStream_t stream) + { + return ForEachCopy(first, last, op, ::cuda::stream_ref{stream}); + } + + /********************************************************************************************************************* + * ForEachInExtents + ********************************************************************************************************************/ + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Iterate through a multi-dimensional extents into a single linear index and a list of indices for each extent + //! dimension. + //! + //! .. versionadded:: 2.4.0 + //! First appears in CUDA Toolkit 12.5. + //! + //! - a single linear index that represents the current iteration + //! - indices of each extent dimension + //! + //! Then apply a function object to the results. + //! + //! - The return value of ``op``, if any, is ignored. + //! + //! **Note**: ``DeviceFor::ForEachInExtents`` supports integral index type up to 64-bits. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use ``ForEachInExtents`` to tabulate a 3D array with its + //! coordinates. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_each_in_extents_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-in-extents-op + //! :end-before: example-end for-each-in-extents-op + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_each_in_extents_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-in-extents-example + //! :end-before: example-end for-each-in-extents-example + //! + //! @endrst + //! + //! @tparam IndexType + //! is an integral type that represents the extent index space (automatically deduced) + //! + //! @tparam Extents + //! are the extent sizes for each rank index (automatically deduced) + //! + //! @tparam OpType + //! is a function object with arity equal to the number of extents + 1 for the linear index (iteration) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] extents + //! Extents object that represents a multi-dimensional index space + //! + //! @param[in] op + //! Function object to apply to each linear index (iteration) and multi-dimensional coordinates + //! + //! @param[in] stream + //! CUDA stream to launch kernels within. Default stream is `NULL` + //! + //! @return cudaError_t + //! error status + template + CUB_RUNTIME_FUNCTION static cudaError_t ForEachInExtents( + void* d_temp_storage, + size_t& temp_storage_bytes, + const ::cuda::std::extents& extents, + OpType op, + cudaStream_t stream = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + return ForEachInExtents(extents, op, stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Iterate through a multi-dimensional extents producing + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! This is an environment-based API that allows customization of: + //! + //! - Stream: Query via ``cuda::get_stream`` + //! + //! - a single linear index that represents the current iteration + //! - list of indices containing the coordinates for each extent dimension + //! + //! Then apply a function object to each tuple of linear index and multidimensional coordinate list. + //! + //! - The return value of ``op``, if any, is ignored. + //! + //! **Note**: ``DeviceFor::ForEachInExtents`` supports integral index type up to 64-bits. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use ``ForEachInExtents`` to tabulate a 3D array with its + //! coordinates. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_each_in_extents_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-in-extents-op + //! :end-before: example-end for-each-in-extents-op + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_each_in_extents_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-in-extents-example + //! :end-before: example-end for-each-in-extents-example + //! + //! @endrst + //! + //! @tparam IndexType + //! is an integral type that represents the extent index space (automatically deduced) + //! + //! @tparam Extents + //! are the extent sizes for each rank index (automatically deduced) + //! + //! @tparam OpType + //! is a function object with arity equal to the number of extents + 1 for the linear index (iteration) + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] extents + //! Extents object that represents a multi-dimensional index space + //! + //! @param[in] op + //! Function object to apply to each linear index (iteration) and multi-dimensional coordinates + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + //! @return cudaError_t + //! error status + template , + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachInExtents(const ::cuda::std::extents& extents, OpType op, const EnvT& env = {}) + { + using extents_type = ::cuda::std::extents; + return cub::DeviceFor::ForEachInLayout(::cuda::std::layout_right::mapping{extents}, op, env); + } + + // We keep this overload around to support types that are convertible to `cudaStream_t` but not copyable + template + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachInExtents(const ::cuda::std::extents& extents, OpType op, cudaStream_t stream) + { + using extents_type = ::cuda::std::extents; + return cub::DeviceFor::ForEachInLayout(::cuda::std::layout_right::mapping{extents}, op, stream); + } + + /********************************************************************************************************************* + * ForEachInLayout + ********************************************************************************************************************/ + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Iterate through multi-dimensional extents using a specific mdspan layout, applying a function object for each + //! element, passing + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - a single linear index that represents the current iteration + //! - a list of indices containing the coordinates for each extent dimension + //! + //! The iteration order depends on the layout type: + //! + //! - ``layout_right``: Iterates in row-major order (rightmost index varies fastest) + //! - ``layout_left``: Iterates in column-major order (leftmost index varies fastest) + //! + //! This is an environment-based API that allows customization of: + //! + //! - Stream: Query via ``cuda::get_stream`` + //! + //! .. note:: + //! + //! The return value of ``op``, if any, is ignored. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The following code snippet demonstrates how to use ``ForEachInLayout`` to iterate through a 2D matrix in + //! column-major order using ``layout_left``. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_each_in_layout_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-in-layout-op + //! :end-before: example-end for-each-in-layout-op + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_for_each_in_layout_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin for-each-in-layout-example + //! :end-before: example-end for-each-in-layout-example + //! + //! @endrst + //! + //! @tparam Layout + //! **[inferred]** The mdspan layout type, must be either ``cuda::std::layout_left`` or ``cuda::std::layout_right`` + //! + //! @tparam IndexType + //! **[inferred]** An integral type that represents the extent index space + //! + //! @tparam Extents + //! **[inferred]** The extent sizes for each rank index + //! + //! @tparam OpType + //! **[inferred]** A function object with arity equal to the number of extents + 1 for the linear index (iteration). + //! The first parameter is the linear index, followed by one parameter for each dimension coordinate. + //! + //! @param[in] layout_mapping + //! Layout mapping object that determines the iteration order and represents a multi-dimensional index space + //! + //! @param[in] op + //! Function object to apply to each linear index (iteration) and multi-dimensional coordinates. + //! Called as ``op(linear_index, coord_0, coord_1, ..., coord_n)`` + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! + //! @return cudaError_t + //! error status + _CCCL_TEMPLATE(typename LayoutMapping, typename OpType, typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES(::cuda::std::__is_cuda_std_layout_left_or_right_mapping_v _CCCL_AND( + !::cuda::std::is_convertible_v)) + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachInLayout(const LayoutMapping& layout_mapping, OpType op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceFor::ForEachInExtents"); + return __for_each_in_extents(layout_mapping, op, env); + } + + // We keep this overload around to support types that are convertible to `cudaStream_t` but not copyable + _CCCL_TEMPLATE(typename LayoutMapping, typename OpType) + _CCCL_REQUIRES(::cuda::std::__is_cuda_std_layout_left_or_right_mapping_v) + CUB_RUNTIME_FUNCTION static cudaError_t + ForEachInLayout(const LayoutMapping& layout_mapping, OpType op, cudaStream_t stream) + { + return ForEachInLayout(layout_mapping, op, ::cuda::stream_ref{stream}); + } + + // Internal version of ForEachInLayout without NVTX range, for use by other device algorithms + _CCCL_TEMPLATE(typename LayoutMapping, typename OpType, typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES(::cuda::std::__is_cuda_std_layout_left_or_right_mapping_v) + CUB_RUNTIME_FUNCTION static cudaError_t + __for_each_in_extents(const LayoutMapping& layout_mapping, OpType op, const EnvT& env = {}) + { + using namespace cub::detail; + using extents_type = typename LayoutMapping::extents_type; + using extent_index_type = typename extents_type::index_type; + using fast_mod_array_t = ::cuda::std::array, extents_type::rank()>; + static constexpr auto seq = ::cuda::std::make_index_sequence{}; + constexpr bool is_layout_right = ::cuda::std::__is_cuda_std_layout_right_mapping_v; + auto extents = layout_mapping.extents(); + fast_mod_array_t sub_sizes_div_array = cub::detail::sub_sizes_fast_div_mod(extents, seq); + fast_mod_array_t extents_div_array = cub::detail::extents_fast_div_mod(extents, seq); + for_each::op_wrapper_extents_t op_wrapper{ + op, extents, sub_sizes_div_array, extents_div_array}; + using ShapeT = implicit_prom_t; + auto shape = static_cast(cub::detail::size(extents)); + if (shape == 0) + { + return cudaSuccess; + } + return __bulk(shape, op_wrapper, env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED + + _CCCL_TEMPLATE(typename LayoutMapping, typename OpType) + _CCCL_REQUIRES(::cuda::std::__is_cuda_std_layout_left_or_right_mapping_v) + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ForEachInLayout( + void* d_temp_storage, + size_t& temp_storage_bytes, + const LayoutMapping& layout_mapping, + OpType op, + cudaStream_t stream = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + return ForEachInLayout(layout_mapping, op, stream); + } + +#endif // !_CCCL_DOXYGEN_INVOKED +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_histogram.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_histogram.cuh new file mode 100644 index 00000000..28a3d45d --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_histogram.cuh @@ -0,0 +1,2573 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceHistogram provides device-wide parallel operations for constructing histogram(s) from a sequence of +//! samples data residing within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceHistogram provides device-wide parallel operations for constructing histogram(s) from a sequence of +//! samples data residing within device-accessible memory. +//! +//! Overview +//! ++++++++++++++++++++++++++ +//! +//! A `histogram `_ counts the number of observations that fall into each +//! of the disjoint categories (known as *bins*). +//! +//! Usage Considerations +//! ++++++++++++++++++++++++++ +//! +//! @cdp_class{DeviceHistogram} +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceHistogram that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::HistogramPolicy`, as shown in the +//! example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin histogram-even-policy-selector +//! :end-before: example-end histogram-even-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin histogram-even-tuning +//! :end-before: example-end histogram-even-tuning +//! +//! @endrst +struct DeviceHistogram +{ + //! @name Evenly-segmented bin ranges + //! @{ + + //! @rst + //! Computes an intensity histogram from a sequence of data samples using equal-width bins. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The number of histogram bins is (``num_levels - 1``) + //! - All bins comprise the same width of sample values: ``(upper_level - lower_level) / (num_levels - 1)``. + //! - If the common type of ``SampleT`` and ``LevelT`` is of integral type, the bin for a sample is + //! computed as ``(sample - lower_level) * (num_levels - 1) / (upper_level - lower_level)``, round + //! down to the nearest whole number. To protect against potential overflows, if the product + //! ``(upper_level - lower_level) * (num_levels - 1)`` exceeds the number representable by an + //! ``uint64_t``, the cuda error ``cudaErrorInvalidValue`` is returned. If the common type is 128 + //! bits wide, bin computation will use 128-bit arithmetic and ``cudaErrorInvalidValue`` will only + //! be returned if bin computation would overflow for 128-bit arithmetic. + //! - The ranges ``[d_samples, d_samples + num_samples)`` and + //! ``[d_histogram, d_histogram + num_levels - 1)`` shall not overlap in any way. + //! - ``cuda::std::common_type`` must be valid, and both LevelT and SampleT must be valid + //! arithmetic types. The common type must be convertible to ``int`` and trivially copyable. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the computation of a six-bin histogram + //! from a sequence of float samples + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input samples and output histogram + //! int num_samples; // e.g., 10 + //! float* d_samples; // e.g., [2.2, 6.1, 7.1, 2.9, 3.5, 0.3, 2.9, 2.1, 6.1, 999.5] + //! int* d_histogram; // e.g., [ -, -, -, -, -, -] + //! int num_levels; // e.g., 7 (seven level boundaries for six bins) + //! float lower_level; // e.g., 0.0 (lower sample value boundary of lowest bin) + //! float upper_level; // e.g., 12.0 (upper sample value boundary of upper bin) + //! ... + //! + //! // Determine temporary device storage requirements + //! void* d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceHistogram::HistogramEven( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, + //! lower_level, upper_level, num_samples); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Compute histograms + //! cub::DeviceHistogram::HistogramEven( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, + //! lower_level, upper_level, num_samples); + //! + //! // d_histogram <-- [1, 5, 0, 3, 0, 0]; + //! + //! @endrst + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_samples + //! The pointer to the input sequence of data samples. + //! + //! @param[out] d_histogram + //! The pointer to the histogram counter output array of length + //! `num_levels - 1`. + //! + //! @param[in] num_levels + //! The number of boundaries (levels) for delineating histogram samples. + //! Implies that the number of bins is `num_levels - 1`. + //! + //! @param[in] lower_level + //! The lower sample value bound (inclusive) for the lowest histogram bin. + //! + //! @param[in] upper_level + //! The upper sample value bound (exclusive) for the highest histogram bin. + //! + //! @param[in] num_samples + //! The number of input samples (i.e., the length of `d_samples`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t HistogramEven( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + CounterT* d_histogram, + int num_levels, + LevelT lower_level, + LevelT upper_level, + OffsetT num_samples, + const EnvT& env = {}) + { + /// The sample value type of the input iterator + using SampleT = cub::detail::it_value_t; + return MultiHistogramEven<1, 1>( + d_temp_storage, + temp_storage_bytes, + d_samples, + ::cuda::std::array{d_histogram}, + ::cuda::std::array{num_levels}, + ::cuda::std::array{lower_level}, + ::cuda::std::array{upper_level}, + num_samples, + static_cast(1), + sizeof(SampleT) * num_samples, + env); + } + + //! @rst + //! Computes an intensity histogram from a sequence of data samples using equal-width bins. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - A two-dimensional *region of interest* within ``d_samples`` can be specified using + //! the ``num_row_samples``, ``num_rows``, and ``row_stride_bytes`` parameters. + //! - The row stride must be a whole multiple of the sample data type + //! size, i.e., ``(row_stride_bytes % sizeof(SampleT)) == 0``. + //! - The number of histogram bins is (``num_levels - 1``) + //! - All bins comprise the same width of sample values: ``(upper_level - lower_level) / (num_levels - 1)`` + //! - If the common type of ``SampleT`` and ``LevelT`` is of integral type, the bin for a sample is + //! computed as ``(sample - lower_level) * (num_levels - 1) / (upper_level - lower_level)``, round + //! down to the nearest whole number. To protect against potential overflows, if the product + //! ``(upper_level - lower_level) * (num_levels - 1)`` exceeds the number representable by an + //! ``uint64_t``, the cuda error ``cudaErrorInvalidValue`` is returned. If the common type is 128 + //! bits wide, bin computation will use 128-bit arithmetic and ``cudaErrorInvalidValue`` will only + //! be returned if bin computation would overflow for 128-bit arithmetic. + //! - For a given row ``r`` in ``[0, num_rows)``, let + //! ``row_begin = d_samples + r * row_stride_bytes / sizeof(SampleT)`` and + //! ``row_end = row_begin + num_row_samples``. The ranges + //! ``[row_begin, row_end)`` and ``[d_histogram, d_histogram + num_levels - 1)`` + //! shall not overlap in any way. + //! - ``cuda::std::common_type`` must be valid, and both LevelT + //! and SampleT must be valid arithmetic types. The common type must be + //! convertible to ``int`` and trivially copyable. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the computation of a six-bin histogram + //! from a 2x5 region of interest within a flattened 2x7 array of float samples. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input samples and output histogram + //! int num_row_samples; // e.g., 5 + //! int num_rows; // e.g., 2; + //! size_t row_stride_bytes; // e.g., 7 * sizeof(float) + //! float* d_samples; // e.g., [2.2, 6.1, 7.1, 2.9, 3.5, -, -, + //! // 0.3, 2.9, 2.1, 6.1, 999.5, -, -] + //! int* d_histogram; // e.g., [ -, -, -, -, -, -] + //! int num_levels; // e.g., 7 (seven level boundaries for six bins) + //! float lower_level; // e.g., 0.0 (lower sample value boundary of lowest bin) + //! float upper_level; // e.g., 12.0 (upper sample value boundary of upper bin) + //! ... + //! + //! // Determine temporary device storage requirements + //! void* d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceHistogram::HistogramEven( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, lower_level, upper_level, + //! num_row_samples, num_rows, row_stride_bytes); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Compute histograms + //! cub::DeviceHistogram::HistogramEven( + //! d_temp_storage, temp_storage_bytes, d_samples, d_histogram, + //! d_samples, d_histogram, num_levels, lower_level, upper_level, + //! num_row_samples, num_rows, row_stride_bytes); + //! + //! // d_histogram <-- [1, 5, 0, 3, 0, 0]; + //! + //! @endrst + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading + //! input samples. @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_samples + //! The pointer to the input sequence of data samples. + //! + //! @param[out] d_histogram + //! The pointer to the histogram counter output array of + //! length `num_levels - 1`. + //! + //! @param[in] num_levels + //! The number of boundaries (levels) for delineating histogram samples. + //! Implies that the number of bins is `num_levels - 1`. + //! + //! @param[in] lower_level + //! The lower sample value bound (inclusive) for the lowest histogram bin. + //! + //! @param[in] upper_level + //! The upper sample value bound (exclusive) for the highest histogram bin. + //! + //! @param[in] num_row_samples + //! The number of data samples per row in the region of interest + //! + //! @param[in] num_rows + //! The number of rows in the region of interest + //! + //! @param[in] row_stride_bytes + //! The number of bytes between starts of consecutive rows in + //! the region of interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t HistogramEven( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + CounterT* d_histogram, + int num_levels, + LevelT lower_level, + LevelT upper_level, + OffsetT num_row_samples, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + return MultiHistogramEven<1, 1>( + d_temp_storage, + temp_storage_bytes, + d_samples, + ::cuda::std::array{d_histogram}, + ::cuda::std::array{num_levels}, + ::cuda::std::array{lower_level}, + ::cuda::std::array{upper_level}, + num_row_samples, + num_rows, + row_stride_bytes, + env); + } + + //! @rst + //! Computes per-channel intensity histograms from a sequence of multi-channel "pixel" data samples using + //! equal-width bins. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The input is a sequence of *pixel* structures, where each pixel comprises + //! a record of ``NUM_CHANNELS`` consecutive data samples + //! (e.g., an *RGBA* pixel). + //! - ``NUM_CHANNELS`` can be up to 4. + //! - Of the ``NUM_CHANNELS`` specified, the function will only compute + //! histograms for the first ``NUM_ACTIVE_CHANNELS`` + //! (e.g., only *RGB* histograms from *RGBA* pixel samples). + //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! - For channel\ :sub:`i`, the range of values for all histogram bins have the same width: + //! ``(upper_level[i] - lower_level[i]) / (num_levels[i] - 1)`` + //! - If the common type of sample and level is of integral type, the bin for a sample is + //! computed as ``(sample - lower_level[i]) * (num_levels - 1) / (upper_level[i] - lower_level[i])``, round down + //! to the nearest whole number. To protect against potential overflows, if, for any channel ``i``, the product + //! ``(upper_level[i] - lower_level[i]) * (num_levels[i] - 1)`` exceeds the number representable by an ``uint64_t``, + //! the cuda error ``cudaErrorInvalidValue`` is returned. If the common type is 128 bits wide, bin computation + //! will use 128-bit arithmetic and ``cudaErrorInvalidValue`` will only be returned if bin + //! computation would overflow for 128-bit arithmetic. + //! - For a given channel ``c`` in ``[0, NUM_ACTIVE_CHANNELS)``, the ranges + //! ``[d_samples, d_samples + NUM_CHANNELS * num_pixels)`` and + //! ``[d_histogram[c], d_histogram[c] + num_levels[c] - 1)`` shall not overlap in any way. + //! - ``cuda::std::common_type`` must be valid, and both LevelT + //! and SampleT must be valid arithmetic types. + //! The common type must be convertible to ``int`` and trivially copyable. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the computation of three 256-bin *RGB* histograms + //! from a quad-channel sequence of *RGBA* pixels (8 bits per channel per pixel) + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input samples and output histograms + //! int num_pixels; // e.g., 5 + //! unsigned char* d_samples; // e.g., [(2, 6, 7, 5), (3, 0, 2, 1), (7, 0, 6, 2), + //! // (0, 6, 7, 5), (3, 0, 2, 6)] + //! int* d_histogram[3]; // e.g., three device pointers to three device buffers, + //! // each allocated with 256 integer counters + //! int num_levels[3]; // e.g., {257, 257, 257}; + //! unsigned int lower_level[3]; // e.g., {0, 0, 0}; + //! unsigned int upper_level[3]; // e.g., {256, 256, 256}; + //! ... + //! + //! // Determine temporary device storage requirements + //! void* d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceHistogram::MultiHistogramEven<4, 3>( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, + //! lower_level, upper_level, num_pixels); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Compute histograms + //! cub::DeviceHistogram::MultiHistogramEven<4, 3>( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, + //! lower_level, upper_level, num_pixels); + //! + //! // d_histogram <-- [ [1, 0, 1, 2, 0, 0, 0, 1, 0, 0, 0, ..., 0], + //! // [0, 3, 0, 0, 0, 0, 2, 0, 0, 0, 0, ..., 0], + //! // [0, 0, 2, 0, 0, 0, 1, 2, 0, 0, 0, ..., 0] ] + //! + //! @endrst + //! + //! @tparam NUM_CHANNELS + //! Number of channels interleaved in the input data (may be greater than + //! the number of channels being actively histogrammed) + //! + //! @tparam NUM_ACTIVE_CHANNELS + //! **[inferred]** Number of channels actively being histogrammed + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading + //! input samples. @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_samples + //! The pointer to the multi-channel input sequence of data samples. + //! The samples from different channels are assumed to be interleaved + //! (e.g., an array of 32-bit pixels where each pixel consists of four + //! *RGBA* 8-bit samples). + //! + //! @param[out] d_histogram + //! @rst + //! The pointers to the histogram counter output arrays, one for each active + //! channel. For channel\ :sub:`i`, the allocation length of + //! ``d_histogram[i]`` should be `num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] num_levels + //! @rst + //! The number of boundaries (levels) for delineating histogram samples in each active channel. + //! Implies that the number of bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] lower_level + //! The lower sample value bound (inclusive) for the lowest histogram bin in each active channel. + //! + //! @param[in] upper_level + //! The upper sample value bound (exclusive) for the highest histogram bin in each active channel. + //! + //! @param[in] num_pixels + //! The number of multi-channel pixels (i.e., the length of `d_samples / NUM_CHANNELS`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t MultiHistogramEven( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_histogram, + ::cuda::std::array num_levels, + ::cuda::std::array lower_level, + ::cuda::std::array upper_level, + OffsetT num_pixels, + const EnvT& env = {}) + { + /// The sample value type of the input iterator + using SampleT = cub::detail::it_value_t; + + return MultiHistogramEven( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_histogram, + num_levels, + lower_level, + upper_level, + num_pixels, + static_cast(1), + sizeof(SampleT) * NUM_CHANNELS * num_pixels, + env); + } + +private: + template + _CCCL_HOST_DEVICE static auto to_array(T* ptr) + { + ::cuda::std::array<::cuda::std::remove_const_t, N> a{}; + ::cuda::std::copy(ptr, ptr + N, a.begin()); + return a; + } + +public: + //! Deprecate [Since 3.0] + template > + CCCL_DEPRECATED_BECAUSE("Prefer the new overload taking cuda::std::arrays") CUB_RUNTIME_FUNCTION static cudaError_t + MultiHistogramEven( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + CounterT* d_histogram[NUM_ACTIVE_CHANNELS], + const int num_levels[NUM_ACTIVE_CHANNELS], + const LevelT lower_level[NUM_ACTIVE_CHANNELS], + const LevelT upper_level[NUM_ACTIVE_CHANNELS], + OffsetT num_pixels, + const EnvT& env = {}) + { + /// The sample value type of the input iterator + using SampleT = cub::detail::it_value_t; + return MultiHistogramEven( + d_temp_storage, + temp_storage_bytes, + d_samples, + to_array(d_histogram), + to_array(num_levels), + to_array(lower_level), + to_array(upper_level), + num_pixels, + env); + } + + //! @rst + //! Computes per-channel intensity histograms from a sequence of + //! multi-channel "pixel" data samples using equal-width bins. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The input is a sequence of *pixel* structures, where each pixel + //! comprises a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel). + //! - ``NUM_CHANNELS`` can be up to 4. + //! - Of the ``NUM_CHANNELS`` specified, the function will only compute + //! histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., only *RGB* + //! histograms from *RGBA* pixel samples). + //! - A two-dimensional *region of interest* within ``d_samples`` can be + //! specified using the ``num_row_samples``, ``num_rows``, and ``row_stride_bytes`` parameters. + //! - The row stride must be a whole multiple of the sample data type + //! size, i.e., ``(row_stride_bytes % sizeof(SampleT)) == 0``. + //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! - For channel\ :sub:`i`, the range of values for all histogram bins have the same width: + //! ``(upper_level[i] - lower_level[i]) / (num_levels[i] - 1)`` + //! - If the common type of sample and level is of integral type, the bin for a sample is + //! computed as ``(sample - lower_level[i]) * (num_levels - 1) / (upper_level[i] - lower_level[i])``, + //! round down to the nearest whole number. To protect against potential overflows, if, for any channel ``i``, + //! the product ``(upper_level[i] - lower_level[i]) * (num_levels[i] - 1)`` exceeds the number representable by + //! an ``uint64_t``, the cuda error ``cudaErrorInvalidValue`` is returned. + //! If the common type is 128 bits wide, bin computation will use 128-bit arithmetic and ``cudaErrorInvalidValue`` + //! will only be returned if bin computation would overflow for 128-bit arithmetic. + //! - For a given row ``r`` in ``[0, num_rows)``, and sample ``s`` in + //! ``[0, num_row_pixels)``, let + //! ``row_begin = d_samples + r * row_stride_bytes / sizeof(SampleT)``, + //! ``sample_begin = row_begin + s * NUM_CHANNELS``, and + //! ``sample_end = sample_begin + NUM_ACTIVE_CHANNELS``. For a given channel ``c`` in + //! ``[0, NUM_ACTIVE_CHANNELS)``, the ranges + //! ``[sample_begin, sample_end)`` and + //! ``[d_histogram[c], d_histogram[c] + num_levels[c] - 1)`` shall not overlap in any way. + //! - ``cuda::std::common_type`` must be valid, and both LevelT + //! and SampleT must be valid arithmetic types. The common type must be + //! convertible to ``int`` and trivially copyable. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the computation of three 256-bin + //! *RGB* histograms from a 2x3 region of interest of within a flattened 2x4 + //! array of quad-channel *RGBA* pixels (8 bits per channel per pixel). + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for input + //! // samples and output histograms + //! int num_row_pixels; // e.g., 3 + //! int num_rows; // e.g., 2 + //! size_t row_stride_bytes; // e.g., 4 * sizeof(unsigned char) * NUM_CHANNELS + //! unsigned char* d_samples; // e.g., [(2, 6, 7, 5), (3, 0, 2, 1), (7, 0, 6, 2), (-, -, -, -), + //! // (0, 6, 7, 5), (3, 0, 2, 6), (1, 1, 1, 1), (-, -, -, -)] + //! int* d_histogram[3]; // e.g., three device pointers to three device buffers, + //! // each allocated with 256 integer counters + //! int num_levels[3]; // e.g., {257, 257, 257}; + //! unsigned int lower_level[3]; // e.g., {0, 0, 0}; + //! unsigned int upper_level[3]; // e.g., {256, 256, 256}; + //! ... + //! + //! // Determine temporary device storage requirements + //! void* d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceHistogram::MultiHistogramEven<4, 3>( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, lower_level, upper_level, + //! num_row_pixels, num_rows, row_stride_bytes); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Compute histograms + //! cub::DeviceHistogram::MultiHistogramEven<4, 3>( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, lower_level, upper_level, + //! num_row_pixels, num_rows, row_stride_bytes); + //! + //! // d_histogram <-- [ [1, 1, 1, 2, 0, 0, 0, 1, 0, 0, 0, ..., 0], + //! // [0, 4, 0, 0, 0, 0, 2, 0, 0, 0, 0, ..., 0], + //! // [0, 1, 2, 0, 0, 0, 1, 2, 0, 0, 0, ..., 0] ] + //! + //! @endrst + //! + //! @tparam NUM_CHANNELS + //! Number of channels interleaved in the input data (may be greater than + //! the number of channels being actively histogrammed) + //! + //! @tparam NUM_ACTIVE_CHANNELS + //! **[inferred]** Number of channels actively being histogrammed + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input + //! samples. @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_samples + //! The pointer to the multi-channel input sequence of data samples. The + //! samples from different channels are assumed to be interleaved (e.g., + //! an array of 32-bit pixels where each pixel consists of four + //! *RGBA* 8-bit samples). + //! + //! @param[out] d_histogram + //! @rst + //! The pointers to the histogram counter output arrays, one for each + //! active channel. For channel\ :sub:`i`, the allocation length + //! of ``d_histogram[i]`` should be ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] num_levels + //! @rst + //! The number of boundaries (levels) for delineating histogram samples in each active channel. + //! Implies that the number of bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] lower_level + //! The lower sample value bound (inclusive) for the lowest histogram bin in each active channel. + //! + //! @param[in] upper_level + //! The upper sample value bound (exclusive) for the highest histogram bin in each active channel. + //! + //! @param[in] num_row_pixels + //! The number of multi-channel pixels per row in the region of interest + //! + //! @param[in] num_rows + //! The number of rows in the region of interest + //! + //! @param[in] row_stride_bytes + //! The number of bytes between starts of consecutive rows in the region of + //! interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t MultiHistogramEven( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_histogram, + ::cuda::std::array num_levels, + ::cuda::std::array lower_level, + ::cuda::std::array upper_level, + OffsetT num_row_pixels, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceHistogram::MultiHistogramEven"); + + using SampleT = cub::detail::it_value_t; + ::cuda::std::bool_constant is_byte_sample; + + using default_policy_selector = + detail::histogram::policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, + temp_storage_bytes, + env, + [&](auto policy_selector, void* storage, size_t& bytes, auto stream) -> cudaError_t { + if constexpr (sizeof(OffsetT) > sizeof(int)) + { + if ((static_cast(num_rows) * row_stride_bytes) < static_cast(INT_MAX)) + { + return detail::histogram::dispatch_even( + storage, + bytes, + d_samples, + d_histogram, + num_levels, + lower_level, + upper_level, + (int) num_row_pixels, + (int) num_rows, + (int) (row_stride_bytes / sizeof(SampleT)), + stream, + is_byte_sample, + policy_selector); + } + } + + return detail::histogram::dispatch_even( + storage, + bytes, + d_samples, + d_histogram, + num_levels, + lower_level, + upper_level, + num_row_pixels, + num_rows, + (OffsetT) (row_stride_bytes / sizeof(SampleT)), + stream, + is_byte_sample, + policy_selector); + }); + } + + //! Deprecate [Since 3.0] + template > + CCCL_DEPRECATED_BECAUSE("Prefer the new overload taking cuda::std::arrays") CUB_RUNTIME_FUNCTION static cudaError_t + MultiHistogramEven( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + CounterT* d_histogram[NUM_ACTIVE_CHANNELS], + const int num_levels[NUM_ACTIVE_CHANNELS], + const LevelT lower_level[NUM_ACTIVE_CHANNELS], + const LevelT upper_level[NUM_ACTIVE_CHANNELS], + OffsetT num_row_pixels, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + return MultiHistogramEven( + d_temp_storage, + temp_storage_bytes, + d_samples, + to_array(d_histogram), + to_array(num_levels), + to_array(lower_level), + to_array(upper_level), + num_row_pixels, + num_rows, + row_stride_bytes, + env); + } + + //! @} + //! @name Custom bin ranges + //! @{ + + //! @rst + //! Computes an intensity histogram from a sequence of data samples using the specified bin boundary levels. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The number of histogram bins is (``num_levels - 1``) + //! - The value range for bin\ :sub:`i` is ``[level[i], level[i+1])`` + //! - The range ``[d_histogram, d_histogram + num_levels - 1)`` shall not + //! overlap ``[d_samples, d_samples + num_samples)`` nor + //! ``[d_levels, d_levels + num_levels)`` in any way. The ranges + //! ``[d_levels, d_levels + num_levels)`` and + //! ``[d_samples, d_samples + num_samples)`` may overlap. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the computation of an six-bin histogram + //! from a sequence of float samples + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for input + //! // samples and output histogram + //! int num_samples; // e.g., 10 + //! float* d_samples; // e.g., [2.2, 6.0, 7.1, 2.9, 3.5, 0.3, 2.9, 2.0, 6.1, 999.5] + //! int* d_histogram; // e.g., [ -, -, -, -, -, -] + //! int num_levels // e.g., 7 (seven level boundaries for six bins) + //! float* d_levels; // e.g., [0.0, 2.0, 4.0, 6.0, 8.0, 12.0, 16.0] + //! ... + //! + //! // Determine temporary device storage requirements + //! void* d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceHistogram::HistogramRange( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, d_levels, num_samples); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Compute histograms + //! cub::DeviceHistogram::HistogramRange( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, d_levels, num_samples); + //! + //! // d_histogram <-- [1, 5, 0, 3, 0, 0]; + //! + //! @endrst + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading + //! input samples. @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_samples + //! The pointer to the input sequence of data samples. + //! + //! @param[out] d_histogram + //! The pointer to the histogram counter output array of length + //! `num_levels - 1`. + //! + //! @param[in] num_levels + //! The number of boundaries (levels) for delineating histogram samples. + //! Implies that the number of bins is `num_levels - 1`. + //! + //! @param[in] d_levels + //! The pointer to the array of boundaries (levels). Bin ranges are defined + //! by consecutive boundary pairings: lower sample value boundaries are + //! inclusive and upper sample value boundaries are exclusive. + //! + //! @param[in] num_samples + //! The number of data samples per row in the region of interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t HistogramRange( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + CounterT* d_histogram, + int num_levels, + const LevelT* d_levels, + OffsetT num_samples, + const EnvT& env = {}) + { + /// The sample value type of the input iterator + using SampleT = cub::detail::it_value_t; + return MultiHistogramRange<1, 1>( + d_temp_storage, + temp_storage_bytes, + d_samples, + ::cuda::std::array{d_histogram}, + ::cuda::std::array{num_levels}, + ::cuda::std::array{d_levels}, + num_samples, + (OffsetT) 1, + (size_t) (sizeof(SampleT) * num_samples), + env); + } + + //! @rst + //! Computes an intensity histogram from a sequence of data samples using the specified bin boundary levels. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - A two-dimensional *region of interest* within ``d_samples`` can be + //! specified using the ``num_row_samples``, ``num_rows``, and ``row_stride_bytes`` parameters. + //! - The row stride must be a whole multiple of the sample data type + //! size, i.e., ``(row_stride_bytes % sizeof(SampleT)) == 0``. + //! - The number of histogram bins is (``num_levels - 1``) + //! - The value range for bin\ :sub:`i` is ``[level[i], level[i+1])`` + //! - For a given row ``r`` in ``[0, num_rows)``, let + //! ``row_begin = d_samples + r * row_stride_bytes / sizeof(SampleT)`` and + //! ``row_end = row_begin + num_row_samples``. The range + //! ``[d_histogram, d_histogram + num_levels - 1)`` shall not overlap + //! ``[row_begin, row_end)`` nor ``[d_levels, d_levels + num_levels)``. + //! The ranges ``[d_levels, d_levels + num_levels)`` and ``[row_begin, row_end)`` may overlap. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the computation of a six-bin histogram + //! from a 2x5 region of interest within a flattened 2x7 array of float samples. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for input samples and + //! // output histogram + //! int num_row_samples; // e.g., 5 + //! int num_rows; // e.g., 2; + //! int row_stride_bytes; // e.g., 7 * sizeof(float) + //! float* d_samples; // e.g., [2.2, 6.0, 7.1, 2.9, 3.5, -, -, + //! // 0.3, 2.9, 2.0, 6.1, 999.5, -, -] + //! int* d_histogram; // e.g., [ -, -, -, -, -, -] + //! int num_levels // e.g., 7 (seven level boundaries for six bins) + //! float *d_levels; // e.g., [0.0, 2.0, 4.0, 6.0, 8.0, 12.0, 16.0] + //! ... + //! + //! // Determine temporary device storage requirements + //! void* d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceHistogram::HistogramRange( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, d_levels, + //! num_row_samples, num_rows, row_stride_bytes); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Compute histograms + //! cub::DeviceHistogram::HistogramRange( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, d_levels, + //! num_row_samples, num_rows, row_stride_bytes); + //! + //! // d_histogram <-- [1, 5, 0, 3, 0, 0]; + //! + //! @endrst + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading + //! input samples. @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_samples + //! The pointer to the input sequence of data samples. + //! + //! @param[out] d_histogram + //! The pointer to the histogram counter output array of length + //! `num_levels - 1`. + //! + //! @param[in] num_levels + //! The number of boundaries (levels) for delineating histogram samples. + //! Implies that the number of bins is `num_levels - 1`. + //! + //! @param[in] d_levels + //! The pointer to the array of boundaries (levels). Bin ranges are defined + //! by consecutive boundary pairings: lower sample value boundaries are + //! inclusive and upper sample value boundaries are exclusive. + //! + //! @param[in] num_row_samples + //! The number of data samples per row in the region of interest + //! + //! @param[in] num_rows + //! The number of rows in the region of interest + //! + //! @param[in] row_stride_bytes + //! The number of bytes between starts of consecutive rows in the region + //! of interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t HistogramRange( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + CounterT* d_histogram, + int num_levels, + const LevelT* d_levels, + OffsetT num_row_samples, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + return MultiHistogramRange<1, 1>( + d_temp_storage, + temp_storage_bytes, + d_samples, + ::cuda::std::array{d_histogram}, + ::cuda::std::array{num_levels}, + ::cuda::std::array{d_levels}, + num_row_samples, + num_rows, + row_stride_bytes, + env); + } + + //! @rst + //! Computes per-channel intensity histograms from a sequence of multi-channel "pixel" data samples + //! using the specified bin boundary levels. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The input is a sequence of *pixel* structures, where each pixel + //! comprises a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel). + //! - ``NUM_CHANNELS`` can be up to 4. + //! - Of the ``NUM_CHANNELS`` specified, the function will only compute + //! histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., *RGB* histograms from *RGBA* pixel samples). + //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! - For channel\ :sub:`i`, the range of values for all histogram bins have the same width: + //! ``(upper_level[i] - lower_level[i]) / (num_levels[i] - 1)`` + //! - For given channels ``c1`` and ``c2`` in ``[0, NUM_ACTIVE_CHANNELS)``, the + //! range ``[d_histogram[c1], d_histogram[c1] + num_levels[c1] - 1)`` shall + //! not overlap ``[d_samples, d_samples + NUM_CHANNELS * num_pixels)`` nor + //! ``[d_levels[c2], d_levels[c2] + num_levels[c2])`` in any way. + //! The ranges ``[d_levels[c2], d_levels[c2] + num_levels[c2])`` and + //! ``[d_samples, d_samples + NUM_CHANNELS * num_pixels)`` may overlap. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the computation of three 4-bin *RGB* + //! histograms from a quad-channel sequence of *RGBA* pixels + //! (8 bits per channel per pixel) + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input samples and output histograms + //! int num_pixels; // e.g., 5 + //! unsigned char *d_samples; // e.g., [(2, 6, 7, 5),(3, 0, 2, 1),(7, 0, 6, 2), + //! // (0, 6, 7, 5),(3, 0, 2, 6)] + //! unsigned int *d_histogram[3]; // e.g., [[ -, -, -, -],[ -, -, -, -],[ -, -, -, -]]; + //! int num_levels[3]; // e.g., {5, 5, 5}; + //! unsigned int *d_levels[3]; // e.g., [ [0, 2, 4, 6, 8], + //! // [0, 2, 4, 6, 8], + //! // [0, 2, 4, 6, 8] ]; + //! ... + //! + //! // Determine temporary device storage requirements + //! void* d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceHistogram::MultiHistogramRange<4, 3>( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, d_levels, num_pixels); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Compute histograms + //! cub::DeviceHistogram::MultiHistogramRange<4, 3>( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, d_levels, num_pixels); + //! + //! // d_histogram <-- [ [1, 3, 0, 1], + //! // [3, 0, 0, 2], + //! // [0, 2, 0, 3] ] + //! + //! @endrst + //! + //! @tparam NUM_CHANNELS + //! Number of channels interleaved in the input data (may be greater than + //! the number of channels being actively histogrammed) + //! + //! @tparam NUM_ACTIVE_CHANNELS + //! **[inferred]** Number of channels actively being histogrammed + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading + //! input samples. @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_samples + //! The pointer to the multi-channel input sequence of data samples. + //! The samples from different channels are assumed to be interleaved (e.g., + //! an array of 32-bit pixels where each pixel consists of four *RGBA* + //! 8-bit samples). + //! + //! @param[out] d_histogram + //! @rst + //! The pointers to the histogram counter output arrays, one for each active + //! channel. For channel\ :sub:`i`, the allocation length of + //! ``d_histogram[i]`` should be ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] num_levels + //! @rst + //! The number of boundaries (levels) for delineating histogram samples in + //! each active channel. Implies that the number of bins for + //! channel\ :sub:`i` is ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] d_levels + //! The pointers to the arrays of boundaries (levels), one for each active + //! channel. Bin ranges are defined by consecutive boundary pairings: lower + //! sample value boundaries are inclusive and upper sample value boundaries + //! are exclusive. + //! + //! @param[in] num_pixels + //! The number of multi-channel pixels (i.e., the length of `d_samples / NUM_CHANNELS`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t MultiHistogramRange( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_histogram, + ::cuda::std::array num_levels, + ::cuda::std::array d_levels, + OffsetT num_pixels, + const EnvT& env = {}) + { + /// The sample value type of the input iterator + using SampleT = cub::detail::it_value_t; + + return MultiHistogramRange( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_histogram, + num_levels, + d_levels, + num_pixels, + (OffsetT) 1, + (size_t) (sizeof(SampleT) * NUM_CHANNELS * num_pixels), + env); + } + + //! Deprecate [Since 3.0] + template + CCCL_DEPRECATED_BECAUSE("Prefer the new overload taking cuda::std::arrays") CUB_RUNTIME_FUNCTION static cudaError_t + MultiHistogramRange( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + CounterT* d_histogram[NUM_ACTIVE_CHANNELS], + const int num_levels[NUM_ACTIVE_CHANNELS], + const LevelT* const d_levels[NUM_ACTIVE_CHANNELS], + OffsetT num_pixels, + cudaStream_t stream = nullptr) + { + return MultiHistogramRange( + d_temp_storage, + temp_storage_bytes, + d_samples, + to_array(d_histogram), + to_array(num_levels), + to_array(d_levels), + num_pixels, + stream); + } + + //! @rst + //! Computes per-channel intensity histograms from a sequence of multi-channel "pixel" data samples using + //! the specified bin boundary levels. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The input is a sequence of *pixel* structures, where each pixel comprises + //! a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel). + //! - ``NUM_CHANNELS`` can be up to 4. + //! - Of the ``NUM_CHANNELS`` specified, the function will only compute + //! histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., *RGB* histograms from *RGBA* pixel samples). + //! - A two-dimensional *region of interest* within ``d_samples`` can be + //! specified using the ``num_row_samples``, ``num_rows``, and ``row_stride_bytes`` parameters. + //! - The row stride must be a whole multiple of the sample data type + //! size, i.e., ``(row_stride_bytes % sizeof(SampleT)) == 0``. + //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! - For channel\ :sub:`i`, the range of values for all histogram bins have the same width: + //! ``(upper_level[i] - lower_level[i]) / (num_levels[i] - 1)`` + //! - For a given row ``r`` in ``[0, num_rows)``, and sample ``s`` in ``[0, num_row_pixels)``, let + //! ``row_begin = d_samples + r * row_stride_bytes / sizeof(SampleT)``, + //! ``sample_begin = row_begin + s * NUM_CHANNELS``, and + //! ``sample_end = sample_begin + NUM_ACTIVE_CHANNELS``. For given channels + //! ``c1`` and ``c2`` in ``[0, NUM_ACTIVE_CHANNELS)``, the range + //! ``[d_histogram[c1], d_histogram[c1] + num_levels[c1] - 1)`` shall not overlap + //! ``[sample_begin, sample_end)`` nor + //! ``[d_levels[c2], d_levels[c2] + num_levels[c2])`` in any way. The ranges + //! ``[d_levels[c2], d_levels[c2] + num_levels[c2])`` and + //! ``[sample_begin, sample_end)`` may overlap. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates the computation of three 4-bin *RGB* + //! histograms from a 2x3 region of interest of within a flattened 2x4 array + //! of quad-channel *RGBA* pixels (8 bits per channel per pixel). + //! + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for input + //! // samples and output histograms + //! int num_row_pixels; // e.g., 3 + //! int num_rows; // e.g., 2 + //! size_t row_stride_bytes; // e.g., 4 * sizeof(unsigned char) * NUM_CHANNELS + //! unsigned char* d_samples; // e.g., [(2, 6, 7, 5),(3, 0, 2, 1),(1, 1, 1, 1),(-, -, -, -), + //! // (7, 0, 6, 2),(0, 6, 7, 5),(3, 0, 2, 6),(-, -, -, -)] + //! int* d_histogram[3]; // e.g., [[ -, -, -, -],[ -, -, -, -],[ -, -, -, -]]; + //! int num_levels[3]; // e.g., {5, 5, 5}; + //! unsigned int* d_levels[3]; // e.g., [ [0, 2, 4, 6, 8], + //! // [0, 2, 4, 6, 8], + //! // [0, 2, 4, 6, 8] ]; + //! ... + //! + //! // Determine temporary device storage requirements + //! void* d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceHistogram::MultiHistogramRange<4, 3>( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, d_levels, + //! num_row_pixels, num_rows, row_stride_bytes); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Compute histograms + //! cub::DeviceHistogram::MultiHistogramRange<4, 3>( + //! d_temp_storage, temp_storage_bytes, + //! d_samples, d_histogram, num_levels, + //! d_levels, num_row_pixels, num_rows, row_stride_bytes); + //! + //! // d_histogram <-- [ [2, 3, 0, 1], + //! // [3, 0, 0, 2], + //! // [1, 2, 0, 3] ] + //! + //! @endrst + //! + //! @tparam NUM_CHANNELS + //! Number of channels interleaved in the input data (may be greater than + //! the number of channels being actively histogrammed) + //! + //! @tparam NUM_ACTIVE_CHANNELS + //! **[inferred]** Number of channels actively being histogrammed + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input + //! samples. @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_samples + //! The pointer to the multi-channel input sequence of data samples. The + //! samples from different channels are assumed to be interleaved (e.g., an + //! array of 32-bit pixels where each pixel consists of four + //! *RGBA* 8-bit samples). + //! + //! @param[out] d_histogram + //! @rst + //! The pointers to the histogram counter output arrays, one for each active + //! channel. For channel\ :sub:`i`, the allocation length of + //! ``d_histogram[i]`` should be ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] num_levels + //! @rst + //! The number of boundaries (levels) for delineating histogram samples in + //! each active channel. Implies that the number of bins for + //! channel\ :sub:`i` is ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] d_levels + //! The pointers to the arrays of boundaries (levels), one for each active + //! channel. Bin ranges are defined by consecutive boundary pairings: lower + //! sample value boundaries are inclusive and upper sample value boundaries + //! are exclusive. + //! + //! @param[in] num_row_pixels + //! The number of multi-channel pixels per row in the region of interest + //! + //! @param[in] num_rows + //! The number of rows in the region of interest + //! + //! @param[in] row_stride_bytes + //! The number of bytes between starts of consecutive rows in the + //! region of interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + CUB_RUNTIME_FUNCTION static cudaError_t MultiHistogramRange( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_histogram, + ::cuda::std::array num_levels, + ::cuda::std::array d_levels, + OffsetT num_row_pixels, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceHistogram::MultiHistogramRange"); + + using SampleT = cub::detail::it_value_t; + ::cuda::std::bool_constant is_byte_sample; + + using default_policy_selector = + detail::histogram::policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, + temp_storage_bytes, + env, + [&](auto policy_selector, void* storage, size_t& bytes, auto stream) -> cudaError_t { + if constexpr (sizeof(OffsetT) > sizeof(int)) + { + if ((static_cast(num_rows) * row_stride_bytes) < static_cast(INT_MAX)) + { + return detail::histogram::dispatch_range( + storage, + bytes, + d_samples, + d_histogram, + num_levels, + d_levels, + (int) num_row_pixels, + (int) num_rows, + (int) (row_stride_bytes / sizeof(SampleT)), + stream, + is_byte_sample, + policy_selector); + } + } + + return detail::histogram::dispatch_range( + storage, + bytes, + d_samples, + d_histogram, + num_levels, + d_levels, + num_row_pixels, + num_rows, + (OffsetT) (row_stride_bytes / sizeof(SampleT)), + stream, + is_byte_sample, + policy_selector); + }); + } + + //! Deprecate [Since 3.0] + template + CCCL_DEPRECATED_BECAUSE("Prefer the new overload taking cuda::std::arrays") CUB_RUNTIME_FUNCTION static cudaError_t + MultiHistogramRange( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + CounterT* d_histogram[NUM_ACTIVE_CHANNELS], + const int num_levels[NUM_ACTIVE_CHANNELS], + const LevelT* const d_levels[NUM_ACTIVE_CHANNELS], + OffsetT num_row_pixels, + OffsetT num_rows, + size_t row_stride_bytes, + cudaStream_t stream = nullptr) + { + return MultiHistogramRange( + d_temp_storage, + temp_storage_bytes, + d_samples, + to_array(d_histogram), + to_array(num_levels), + to_array(d_levels), + num_row_pixels, + num_rows, + row_stride_bytes, + stream); + } + + //@} + + //! @name Environment-based overloads + //! @{ + + //! @rst + //! Computes an intensity histogram from a sequence of data samples using equal-width bins. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The number of histogram bins is (``num_levels - 1``) + //! - All bins comprise the same width of sample values: ``(upper_level - lower_level) / (num_levels - 1)``. + //! - If the common type of ``SampleT`` and ``LevelT`` is of integral type, the bin for a sample is + //! computed as ``(sample - lower_level) * (num_levels - 1) / (upper_level - lower_level)``, round + //! down to the nearest whole number. To protect against potential overflows, if the product + //! ``(upper_level - lower_level) * (num_levels - 1)`` exceeds the number representable by an + //! ``uint64_t``, the cuda error ``cudaErrorInvalidValue`` is returned. If the common type is 128 + //! bits wide, bin computation will use 128-bit arithmetic and ``cudaErrorInvalidValue`` will only + //! be returned if bin computation would overflow for 128-bit arithmetic. + //! - The ranges ``[d_samples, d_samples + num_samples)`` and + //! ``[d_histogram, d_histogram + num_levels - 1)`` shall not overlap in any way. + //! - ``cuda::std::common_type`` must be valid, and both LevelT and SampleT must be valid + //! arithmetic types. The common type must be convertible to ``int`` and trivially copyable. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin histogram-even-env + //! :end-before: example-end histogram-even-env + //! + //! @endrst + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, pointer differences, etc. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_samples + //! The pointer to the input sequence of data samples. + //! + //! @param[out] d_histogram + //! The pointer to the histogram counter output array of length `num_levels - 1`. + //! + //! @param[in] num_levels + //! The number of boundaries (levels) for delineating histogram samples. + //! Implies that the number of bins is `num_levels - 1`. + //! + //! @param[in] lower_level + //! The lower sample value bound (inclusive) for the lowest histogram bin. + //! + //! @param[in] upper_level + //! The upper sample value bound (exclusive) for the highest histogram bin. + //! + //! @param[in] num_samples + //! The number of input samples (i.e., the length of `d_samples`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t HistogramEven( + SampleIteratorT d_samples, + CounterT* d_histogram, + int num_levels, + LevelT lower_level, + LevelT upper_level, + OffsetT num_samples, + const EnvT& env = {}) + { + using SampleT = cub::detail::it_value_t; + return MultiHistogramEven<1, 1>( + d_samples, + ::cuda::std::array{d_histogram}, + ::cuda::std::array{num_levels}, + ::cuda::std::array{lower_level}, + ::cuda::std::array{upper_level}, + num_samples, + static_cast(1), + sizeof(SampleT) * num_samples, + env); + } + + //! @rst + //! Computes an intensity histogram from a 2D region of data samples using equal-width bins. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - A two-dimensional *region of interest* within ``d_samples`` can be specified using + //! the ``num_row_samples``, ``num_rows``, and ``row_stride_bytes`` parameters. + //! - The row stride must be a whole multiple of the sample data type + //! size, i.e., ``(row_stride_bytes % sizeof(SampleT)) == 0``. + //! - The number of histogram bins is (``num_levels - 1``) + //! - All bins comprise the same width of sample values: ``(upper_level - lower_level) / (num_levels - 1)`` + //! - If the common type of ``SampleT`` and ``LevelT`` is of integral type, the bin for a sample is + //! computed as ``(sample - lower_level) * (num_levels - 1) / (upper_level - lower_level)``, round + //! down to the nearest whole number. To protect against potential overflows, if the product + //! ``(upper_level - lower_level) * (num_levels - 1)`` exceeds the number representable by an + //! ``uint64_t``, the cuda error ``cudaErrorInvalidValue`` is returned. If the common type is 128 + //! bits wide, bin computation will use 128-bit arithmetic and ``cudaErrorInvalidValue`` will only + //! be returned if bin computation would overflow for 128-bit arithmetic. + //! - For a given row ``r`` in ``[0, num_rows)``, let + //! ``row_begin = d_samples + r * row_stride_bytes / sizeof(SampleT)`` and + //! ``row_end = row_begin + num_row_samples``. The ranges + //! ``[row_begin, row_end)`` and ``[d_histogram, d_histogram + num_levels - 1)`` + //! shall not overlap in any way. + //! - ``cuda::std::common_type`` must be valid, and both LevelT + //! and SampleT must be valid arithmetic types. The common type must be + //! convertible to ``int`` and trivially copyable. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin histogram-even-2d-env + //! :end-before: example-end histogram-even-2d-env + //! + //! @endrst + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, pointer differences, etc. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_samples + //! The pointer to the input sequence of data samples. + //! + //! @param[out] d_histogram + //! The pointer to the histogram counter output array of length `num_levels - 1`. + //! + //! @param[in] num_levels + //! The number of boundaries (levels) for delineating histogram samples. + //! Implies that the number of bins is `num_levels - 1`. + //! + //! @param[in] lower_level + //! The lower sample value bound (inclusive) for the lowest histogram bin. + //! + //! @param[in] upper_level + //! The upper sample value bound (exclusive) for the highest histogram bin. + //! + //! @param[in] num_row_samples + //! The number of data samples per row in the region of interest + //! + //! @param[in] num_rows + //! The number of rows in the region of interest + //! + //! @param[in] row_stride_bytes + //! The number of bytes between starts of consecutive rows in the region of interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t HistogramEven( + SampleIteratorT d_samples, + CounterT* d_histogram, + int num_levels, + LevelT lower_level, + LevelT upper_level, + OffsetT num_row_samples, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + return MultiHistogramEven<1, 1>( + d_samples, + ::cuda::std::array{d_histogram}, + ::cuda::std::array{num_levels}, + ::cuda::std::array{lower_level}, + ::cuda::std::array{upper_level}, + num_row_samples, + num_rows, + row_stride_bytes, + env); + } + + //! @rst + //! Computes per-channel intensity histograms from a sequence of multi-channel "pixel" data samples + //! using equal-width bins. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The input is a sequence of *pixel* structures, where each pixel comprises + //! a record of ``NUM_CHANNELS`` consecutive data samples + //! (e.g., an *RGBA* pixel). + //! - ``NUM_CHANNELS`` can be up to 4. + //! - Of the ``NUM_CHANNELS`` specified, the function will only compute + //! histograms for the first ``NUM_ACTIVE_CHANNELS`` + //! (e.g., only *RGB* histograms from *RGBA* pixel samples). + //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! - For channel\ :sub:`i`, the range of values for all histogram bins have the same width: + //! ``(upper_level[i] - lower_level[i]) / (num_levels[i] - 1)`` + //! - If the common type of sample and level is of integral type, the bin for a sample is + //! computed as ``(sample - lower_level[i]) * (num_levels - 1) / (upper_level[i] - lower_level[i])``, round down + //! to the nearest whole number. To protect against potential overflows, if, for any channel ``i``, the product + //! ``(upper_level[i] - lower_level[i]) * (num_levels[i] - 1)`` exceeds the number representable by an ``uint64_t``, + //! the cuda error ``cudaErrorInvalidValue`` is returned. If the common type is 128 bits wide, bin computation + //! will use 128-bit arithmetic and ``cudaErrorInvalidValue`` will only be returned if bin + //! computation would overflow for 128-bit arithmetic. + //! - For a given channel ``c`` in ``[0, NUM_ACTIVE_CHANNELS)``, the ranges + //! ``[d_samples, d_samples + NUM_CHANNELS * num_pixels)`` and + //! ``[d_histogram[c], d_histogram[c] + num_levels[c] - 1)`` shall not overlap in any way. + //! - ``cuda::std::common_type`` must be valid, and both LevelT + //! and SampleT must be valid arithmetic types. + //! The common type must be convertible to ``int`` and trivially copyable. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin multi-histogram-even-1d-env + //! :end-before: example-end multi-histogram-even-1d-env + //! + //! @endrst + //! + //! @tparam NUM_CHANNELS + //! Number of channels interleaved in the input data (may be greater than the number of channels being + //! actively histogrammed) + //! + //! @tparam NUM_ACTIVE_CHANNELS + //! **[inferred]** Number of channels actively being histogrammed + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, pointer differences, etc. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_samples + //! The pointer to the multi-channel input sequence of data samples. + //! + //! @param[out] d_histogram + //! Array of active channel histogram counter output arrays, each of length `num_levels[channel] - 1`. + //! + //! @param[in] num_levels + //! Array of the number of boundaries (levels) for each active channel. + //! + //! @param[in] lower_level + //! Array of the lower sample value bound (inclusive) for the lowest bin of each active channel. + //! + //! @param[in] upper_level + //! Array of the upper sample value bound (exclusive) for the highest bin of each active channel. + //! + //! @param[in] num_pixels + //! The number of multi-channel pixels (i.e., the length of `d_samples / NUM_CHANNELS`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MultiHistogramEven( + SampleIteratorT d_samples, + ::cuda::std::array d_histogram, + ::cuda::std::array num_levels, + ::cuda::std::array lower_level, + ::cuda::std::array upper_level, + OffsetT num_pixels, + const EnvT& env = {}) + { + using SampleT = cub::detail::it_value_t; + return MultiHistogramEven( + d_samples, + d_histogram, + num_levels, + lower_level, + upper_level, + num_pixels, + static_cast(1), + sizeof(SampleT) * NUM_CHANNELS * num_pixels, + env); + } + + //! @rst + //! Computes per-channel intensity histograms from a 2D region of multi-channel "pixel" data samples + //! using equal-width bins. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The input is a sequence of *pixel* structures, where each pixel + //! comprises a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel). + //! - ``NUM_CHANNELS`` can be up to 4. + //! - Of the ``NUM_CHANNELS`` specified, the function will only compute + //! histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., only *RGB* + //! histograms from *RGBA* pixel samples). + //! - A two-dimensional *region of interest* within ``d_samples`` can be + //! specified using the ``num_row_samples``, ``num_rows``, and ``row_stride_bytes`` parameters. + //! - The row stride must be a whole multiple of the sample data type + //! size, i.e., ``(row_stride_bytes % sizeof(SampleT)) == 0``. + //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! - For channel\ :sub:`i`, the range of values for all histogram bins have the same width: + //! ``(upper_level[i] - lower_level[i]) / (num_levels[i] - 1)`` + //! - If the common type of sample and level is of integral type, the bin for a sample is + //! computed as ``(sample - lower_level[i]) * (num_levels - 1) / (upper_level[i] - lower_level[i])``, + //! round down to the nearest whole number. To protect against potential overflows, if, for any channel ``i``, + //! the product ``(upper_level[i] - lower_level[i]) * (num_levels[i] - 1)`` exceeds the number representable by + //! an ``uint64_t``, the cuda error ``cudaErrorInvalidValue`` is returned. + //! If the common type is 128 bits wide, bin computation will use 128-bit arithmetic and ``cudaErrorInvalidValue`` + //! will only be returned if bin computation would overflow for 128-bit arithmetic. + //! - For a given row ``r`` in ``[0, num_rows)``, and sample ``s`` in + //! ``[0, num_row_pixels)``, let + //! ``row_begin = d_samples + r * row_stride_bytes / sizeof(SampleT)``, + //! ``sample_begin = row_begin + s * NUM_CHANNELS``, and + //! ``sample_end = sample_begin + NUM_ACTIVE_CHANNELS``. For a given channel ``c`` in + //! ``[0, NUM_ACTIVE_CHANNELS)``, the ranges + //! ``[sample_begin, sample_end)`` and + //! ``[d_histogram[c], d_histogram[c] + num_levels[c] - 1)`` shall not overlap in any way. + //! - ``cuda::std::common_type`` must be valid, and both LevelT + //! and SampleT must be valid arithmetic types. The common type must be + //! convertible to ``int`` and trivially copyable. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin multi-histogram-even-2d-env + //! :end-before: example-end multi-histogram-even-2d-env + //! + //! @endrst + //! + //! @tparam NUM_CHANNELS + //! Number of channels interleaved in the input data (may be greater than the number of channels being + //! actively histogrammed) + //! + //! @tparam NUM_ACTIVE_CHANNELS + //! **[inferred]** Number of channels actively being histogrammed + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, + //! pointer differences, etc. @offset_size1 + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_samples + //! The pointer to the multi-channel input sequence of data samples. The + //! samples from different channels are assumed to be interleaved (e.g., + //! an array of 32-bit pixels where each pixel consists of four + //! *RGBA* 8-bit samples). + //! + //! @param[out] d_histogram + //! @rst + //! The pointers to the histogram counter output arrays, one for each + //! active channel. For channel\ :sub:`i`, the allocation length + //! of ``d_histogram[i]`` should be ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] num_levels + //! @rst + //! The number of boundaries (levels) for delineating histogram samples in each active channel. + //! Implies that the number of bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! @endrst + //! + //! @param[in] lower_level + //! The lower sample value bound (inclusive) for the lowest histogram bin in each active channel. + //! + //! @param[in] upper_level + //! The upper sample value bound (exclusive) for the highest histogram bin in each active channel. + //! + //! @param[in] num_row_pixels + //! The number of multi-channel pixels per row in the region of interest + //! + //! @param[in] num_rows + //! The number of rows in the region of interest + //! + //! @param[in] row_stride_bytes + //! The number of bytes between starts of consecutive rows in the region of + //! interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MultiHistogramEven( + SampleIteratorT d_samples, + ::cuda::std::array d_histogram, + ::cuda::std::array num_levels, + ::cuda::std::array lower_level, + ::cuda::std::array upper_level, + OffsetT num_row_pixels, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceHistogram::MultiHistogramEven"); + + using SampleT = cub::detail::it_value_t; + ::cuda::std::bool_constant is_byte_sample; + + using default_policy_selector = + detail::histogram::policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) -> cudaError_t { + if constexpr (sizeof(OffsetT) > sizeof(int)) + { + if ((unsigned long long) (num_rows * row_stride_bytes) < (unsigned long long) INT_MAX) + { + return detail::histogram::dispatch_even( + storage, + bytes, + d_samples, + d_histogram, + num_levels, + lower_level, + upper_level, + (int) num_row_pixels, + (int) num_rows, + (int) (row_stride_bytes / sizeof(SampleT)), + stream, + is_byte_sample, + policy_selector); + } + } + + return detail::histogram::dispatch_even( + storage, + bytes, + d_samples, + d_histogram, + num_levels, + lower_level, + upper_level, + num_row_pixels, + num_rows, + (OffsetT) (row_stride_bytes / sizeof(SampleT)), + stream, + is_byte_sample, + policy_selector); + }); + } + + //! @rst + //! Computes an intensity histogram from a sequence of data samples using the specified bin boundary levels. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The number of histogram bins is (``num_levels - 1``) + //! - The value range for bin\ :sub:`i` is ``[level[i], level[i+1])`` + //! - The range ``[d_histogram, d_histogram + num_levels - 1)`` shall not + //! overlap ``[d_samples, d_samples + num_samples)`` nor + //! ``[d_levels, d_levels + num_levels)`` in any way. The ranges + //! ``[d_levels, d_levels + num_levels)`` and + //! ``[d_samples, d_samples + num_samples)`` may overlap. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin histogram-range-env + //! :end-before: example-end histogram-range-env + //! + //! @endrst + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, pointer differences, etc. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_samples + //! The pointer to the input sequence of data samples. + //! + //! @param[out] d_histogram + //! The pointer to the histogram counter output array of length `num_levels - 1`. + //! + //! @param[in] num_levels + //! The number of boundaries (levels) for delineating histogram samples. + //! Implies that the number of bins is `num_levels - 1`. + //! + //! @param[in] d_levels + //! The pointer to the array of boundaries (levels). Bin ranges are defined + //! by consecutive boundary pairings: lower sample value boundaries are + //! inclusive and upper sample value boundaries are exclusive. + //! + //! @param[in] num_samples + //! The number of input samples (i.e., the length of `d_samples`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t HistogramRange( + SampleIteratorT d_samples, + CounterT* d_histogram, + int num_levels, + const LevelT* d_levels, + OffsetT num_samples, + const EnvT& env = {}) + { + using SampleT = cub::detail::it_value_t; + return MultiHistogramRange<1, 1>( + d_samples, + ::cuda::std::array{d_histogram}, + ::cuda::std::array{num_levels}, + ::cuda::std::array{d_levels}, + num_samples, + static_cast(1), + sizeof(SampleT) * num_samples, + env); + } + + //! @rst + //! Computes an intensity histogram from a 2D region of data samples using the specified bin boundary levels. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - A two-dimensional *region of interest* within ``d_samples`` can be + //! specified using the ``num_row_samples``, ``num_rows``, and ``row_stride_bytes`` parameters. + //! - The row stride must be a whole multiple of the sample data type + //! size, i.e., ``(row_stride_bytes % sizeof(SampleT)) == 0``. + //! - The number of histogram bins is (``num_levels - 1``) + //! - The value range for bin\ :sub:`i` is ``[level[i], level[i+1])`` + //! - For a given row ``r`` in ``[0, num_rows)``, let + //! ``row_begin = d_samples + r * row_stride_bytes / sizeof(SampleT)`` and + //! ``row_end = row_begin + num_row_samples``. The range + //! ``[d_histogram, d_histogram + num_levels - 1)`` shall not overlap + //! ``[row_begin, row_end)`` nor ``[d_levels, d_levels + num_levels)``. + //! The ranges ``[d_levels, d_levels + num_levels)`` and ``[row_begin, row_end)`` may overlap. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin histogram-range-2d-env + //! :end-before: example-end histogram-range-2d-env + //! + //! @endrst + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, pointer differences, etc. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_samples + //! The pointer to the input sequence of data samples. + //! + //! @param[out] d_histogram + //! The pointer to the histogram counter output array of length `num_levels - 1`. + //! + //! @param[in] num_levels + //! The number of boundaries (levels) for delineating histogram samples. + //! Implies that the number of bins is `num_levels - 1`. + //! + //! @param[in] d_levels + //! The pointer to the array of boundaries (levels). Bin ranges are defined + //! by consecutive boundary pairings: lower sample value boundaries are + //! inclusive and upper sample value boundaries are exclusive. + //! + //! @param[in] num_row_samples + //! The number of data samples per row in the region of interest + //! + //! @param[in] num_rows + //! The number of rows in the region of interest + //! + //! @param[in] row_stride_bytes + //! The number of bytes between starts of consecutive rows in the region of interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t HistogramRange( + SampleIteratorT d_samples, + CounterT* d_histogram, + int num_levels, + const LevelT* d_levels, + OffsetT num_row_samples, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + return MultiHistogramRange<1, 1>( + d_samples, + ::cuda::std::array{d_histogram}, + ::cuda::std::array{num_levels}, + ::cuda::std::array{d_levels}, + num_row_samples, + num_rows, + row_stride_bytes, + env); + } + + //! @rst + //! Computes per-channel intensity histograms from a sequence of multi-channel "pixel" data samples + //! using the specified bin boundary levels. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The input is a sequence of *pixel* structures, where each pixel + //! comprises a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel). + //! - ``NUM_CHANNELS`` can be up to 4. + //! - Of the ``NUM_CHANNELS`` specified, the function will only compute + //! histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., *RGB* histograms from *RGBA* pixel samples). + //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! - For channel\ :sub:`i`, the range of values for all histogram bins have the same width: + //! ``(upper_level[i] - lower_level[i]) / (num_levels[i] - 1)`` + //! - For given channels ``c1`` and ``c2`` in ``[0, NUM_ACTIVE_CHANNELS)``, the + //! range ``[d_histogram[c1], d_histogram[c1] + num_levels[c1] - 1)`` shall + //! not overlap ``[d_samples, d_samples + NUM_CHANNELS * num_pixels)`` nor + //! ``[d_levels[c2], d_levels[c2] + num_levels[c2])`` in any way. + //! The ranges ``[d_levels[c2], d_levels[c2] + num_levels[c2])`` and + //! ``[d_samples, d_samples + NUM_CHANNELS * num_pixels)`` may overlap. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin multi-histogram-range-1d-env + //! :end-before: example-end multi-histogram-range-1d-env + //! + //! @endrst + //! + //! @tparam NUM_CHANNELS + //! Number of channels interleaved in the input data (may be greater than the number of channels being + //! actively histogrammed) + //! + //! @tparam NUM_ACTIVE_CHANNELS + //! **[inferred]** Number of channels actively being histogrammed + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, pointer differences, etc. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_samples + //! The pointer to the multi-channel input sequence of data samples. + //! + //! @param[out] d_histogram + //! Array of active channel histogram counter output arrays, each of length `num_levels[channel] - 1`. + //! + //! @param[in] num_levels + //! Array of the number of boundaries (levels) for each active channel. + //! + //! @param[in] d_levels + //! Array of pointers to the arrays of boundaries (levels) for each active channel. + //! + //! @param[in] num_pixels + //! The number of multi-channel pixels (i.e., the length of `d_samples / NUM_CHANNELS`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MultiHistogramRange( + SampleIteratorT d_samples, + ::cuda::std::array d_histogram, + ::cuda::std::array num_levels, + ::cuda::std::array d_levels, + OffsetT num_pixels, + const EnvT& env = {}) + { + using SampleT = cub::detail::it_value_t; + return MultiHistogramRange( + d_samples, + d_histogram, + num_levels, + d_levels, + num_pixels, + static_cast(1), + sizeof(SampleT) * NUM_CHANNELS * num_pixels, + env); + } + + //! @rst + //! Computes per-channel intensity histograms from a 2D region of multi-channel "pixel" data samples + //! using the specified bin boundary levels. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The input is a sequence of *pixel* structures, where each pixel comprises + //! a record of ``NUM_CHANNELS`` consecutive data samples (e.g., an *RGBA* pixel). + //! - ``NUM_CHANNELS`` can be up to 4. + //! - Of the ``NUM_CHANNELS`` specified, the function will only compute + //! histograms for the first ``NUM_ACTIVE_CHANNELS`` (e.g., *RGB* histograms from *RGBA* pixel samples). + //! - A two-dimensional *region of interest* within ``d_samples`` can be + //! specified using the ``num_row_samples``, ``num_rows``, and ``row_stride_bytes`` parameters. + //! - The row stride must be a whole multiple of the sample data type + //! size, i.e., ``(row_stride_bytes % sizeof(SampleT)) == 0``. + //! - The number of histogram bins for channel\ :sub:`i` is ``num_levels[i] - 1``. + //! - For channel\ :sub:`i`, the range of values for all histogram bins have the same width: + //! ``(upper_level[i] - lower_level[i]) / (num_levels[i] - 1)`` + //! - For a given row ``r`` in ``[0, num_rows)``, and sample ``s`` in ``[0, num_row_pixels)``, let + //! ``row_begin = d_samples + r * row_stride_bytes / sizeof(SampleT)``, + //! ``sample_begin = row_begin + s * NUM_CHANNELS``, and + //! ``sample_end = sample_begin + NUM_ACTIVE_CHANNELS``. For given channels + //! ``c1`` and ``c2`` in ``[0, NUM_ACTIVE_CHANNELS)``, the range + //! ``[d_histogram[c1], d_histogram[c1] + num_levels[c1] - 1)`` shall not overlap + //! ``[sample_begin, sample_end)`` nor + //! ``[d_levels[c2], d_levels[c2] + num_levels[c2])`` in any way. The ranges + //! ``[d_levels[c2], d_levels[c2] + num_levels[c2])`` and + //! ``[sample_begin, sample_end)`` may overlap. + //! - @devicestorage + //! + //! Snippet + //! +++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_histogram_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin multi-histogram-range-2d-env + //! :end-before: example-end multi-histogram-range-2d-env + //! + //! @endrst + //! + //! @tparam NUM_CHANNELS + //! Number of channels interleaved in the input data (may be greater than the number of channels being + //! actively histogrammed) + //! + //! @tparam NUM_ACTIVE_CHANNELS + //! **[inferred]** Number of channels actively being histogrammed + //! + //! @tparam SampleIteratorT + //! **[inferred]** Random-access input iterator type for reading input samples @iterator + //! + //! @tparam CounterT + //! **[inferred]** Integer type for histogram bin counters + //! + //! @tparam LevelT + //! **[inferred]** Type for specifying boundaries (levels) + //! + //! @tparam OffsetT + //! **[inferred]** Signed integer type for sequence offsets, list lengths, pointer differences, etc. + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_samples + //! The pointer to the multi-channel input sequence of data samples. + //! + //! @param[out] d_histogram + //! Array of active channel histogram counter output arrays, each of length `num_levels[channel] - 1`. + //! + //! @param[in] num_levels + //! Array of the number of boundaries (levels) for each active channel. + //! + //! @param[in] d_levels + //! Array of pointers to the arrays of boundaries (levels) for each active channel. + //! + //! @param[in] num_row_pixels + //! The number of multi-channel pixels per row in the region of interest + //! + //! @param[in] num_rows + //! The number of rows in the region of interest + //! + //! @param[in] row_stride_bytes + //! The number of bytes between starts of consecutive rows in the region of interest + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MultiHistogramRange( + SampleIteratorT d_samples, + ::cuda::std::array d_histogram, + ::cuda::std::array num_levels, + ::cuda::std::array d_levels, + OffsetT num_row_pixels, + OffsetT num_rows, + size_t row_stride_bytes, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceHistogram::MultiHistogramRange"); + + using SampleT = cub::detail::it_value_t; + ::cuda::std::bool_constant is_byte_sample; + + using default_policy_selector = + detail::histogram::policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) -> cudaError_t { + if constexpr (sizeof(OffsetT) > sizeof(int)) + { + if ((unsigned long long) (num_rows * row_stride_bytes) < (unsigned long long) INT_MAX) + { + return detail::histogram::dispatch_range( + storage, + bytes, + d_samples, + d_histogram, + num_levels, + d_levels, + (int) num_row_pixels, + (int) num_rows, + (int) (row_stride_bytes / sizeof(SampleT)), + stream, + is_byte_sample, + policy_selector); + } + } + + return detail::histogram::dispatch_range( + storage, + bytes, + d_samples, + d_histogram, + num_levels, + d_levels, + num_row_pixels, + num_rows, + (OffsetT) (row_stride_bytes / sizeof(SampleT)), + stream, + is_byte_sample, + policy_selector); + }); + } + + //@} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_memcpy.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_memcpy.cuh new file mode 100644 index 00000000..a372f841 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_memcpy.cuh @@ -0,0 +1,293 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceMemcpy provides device-wide, parallel operations for copying data. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @brief cub::DeviceMemcpy provides device-wide, parallel operations for copying data. +//! +//! @rst +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceMemcpy that accept an environment can be tuned by passing a custom :ref:`policy selector +//! ` that returns a :cpp:struct:`cub::BatchedCopyPolicy`, as shown in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_memcpy_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin memcpy-batched-policy-selector +//! :end-before: example-end memcpy-batched-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_memcpy_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin memcpy-batched-tuning +//! :end-before: example-end memcpy-batched-tuning +//! @endrst +struct DeviceMemcpy +{ + //! @rst + //! Copies data from a batch of given source buffers to their corresponding destination buffer. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! .. note:: + //! + //! If any input buffer aliases memory from any output buffer the behavior is undefined. + //! If any output buffer aliases memory of another output buffer the behavior is undefined. + //! Input buffers can alias one another. + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates usage of DeviceMemcpy::Batched for mutating strings withing + //! a single string buffer. + //! + //! .. code-block:: c++ + //! + //! struct GetPtrToStringItem + //! { + //! __host__ __device__ __forceinline__ void *operator()(uint32_t index) + //! { + //! return &d_string_data_in[d_string_offsets[index]]; + //! } + //! char *d_string_data_in; + //! uint32_t *d_string_offsets; + //! }; + //! + //! struct GetStringItemSize + //! { + //! __host__ __device__ __forceinline__ uint32_t operator()(uint32_t index) + //! { + //! return d_string_offsets[index + 1] - d_string_offsets[index]; + //! } + //! uint32_t *d_string_offsets; + //! }; + //! + //! uint32_t num_strings = 5; + //! char *d_string_data_in; // e.g., "TomatoesBananasApplesOrangesGrapes" + //! char *d_string_data_out; // e.g., " ... " + //! uint32_t *d_string_offsets_old; // e.g., [0, 8, 15, 21, 28, 34] + //! uint32_t *d_string_offsets_new; // e.g., [0, 6, 13, 19, 26, 34] + //! uint32_t *d_gather_index; // e.g., [2, 1, 4, 3, 0] + //! + //! // Initialize an iterator that returns d_gather_index[i] when the i-th item is dereferenced + //! auto gather_iterator = thrust::make_permutation_iterator(thrust::make_counting_iterator(0), + //! d_gather_index); + //! + //! // Returns pointers to the input buffer for each string + //! auto str_ptrs_in = thrust::make_transform_iterator(gather_iterator, + //! GetPtrToStringItem{d_string_data_in, + //! d_string_offsets_old}); + //! + //! // Returns the string size of the i-th string + //! auto str_sizes = thrust::make_transform_iterator(gather_iterator, + //! GetStringItemSize{d_string_offsets_old}); + //! + //! // Returns pointers to the output buffer for each string + //! auto str_ptrs_out = thrust::make_transform_iterator(thrust::make_counting_iterator(0), + //! GetPtrToStringItem{d_string_data_out, + //! d_string_offsets_new}); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceMemcpy::Batched(d_temp_storage, temp_storage_bytes, str_ptrs_in, str_ptrs_out, + //! str_sizes, num_strings); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run batched copy algorithm (used to permute strings) + //! cub::DeviceMemcpy::Batched(d_temp_storage, temp_storage_bytes, str_ptrs_in, str_ptrs_out, + //! str_sizes, num_strings); + //! + //! // d_string_data_out <-- "ApplesBananasGrapesOrangesTomatoe" + //! + //! @endrst + //! + //! @tparam InputBufferIt + //! **[inferred]** Device-accessible random-access input iterator type providing the pointers to + //! the source memory buffers + //! + //! @tparam OutputBufferIt + //! **[inferred]** Device-accessible random-access input iterator type providing the pointers to + //! the destination memory buffers + //! + //! @tparam BufferSizeIteratorT + //! **[inferred]** Device-accessible random-access input iterator type providing the number of bytes + //! to be copied for each pair of buffers + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in] input_buffer_it + //! Device-accessible iterator providing the pointers to the source memory buffers + //! + //! @param[in] output_buffer_it + //! Device-accessible iterator providing the pointers to the destination memory buffers + //! + //! @param[in] buffer_sizes + //! Device-accessible iterator providing the number of bytes to be copied for each pair of buffers + //! + //! @param[in] num_buffers + //! The total number of buffer pairs + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t Batched( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputBufferIt input_buffer_it, + OutputBufferIt output_buffer_it, + BufferSizeIteratorT buffer_sizes, + ::cuda::std::int64_t num_buffers, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceMemcpy::Batched"); + static_assert(::cuda::std::is_pointer_v>, + "DeviceMemcpy::Batched only supports copying of memory buffers." + "Please consider using DeviceCopy::Batched instead."); + static_assert(::cuda::std::is_pointer_v>, + "DeviceMemcpy::Batched only supports copying of memory buffers." + "Please consider using DeviceCopy::Batched instead."); + + // Integer type large enough to hold any offset in [0, num_thread_blocks_launched), where a safe + // upper bound on num_thread_blocks_launched can be assumed to be given by + // IDIV_CEIL(num_buffers, 64) + using BlockOffsetT = uint32_t; + + return detail::batch_memcpy::dispatch( + d_temp_storage, temp_storage_bytes, input_buffer_it, output_buffer_it, buffer_sizes, num_buffers, stream); + } + + //! @rst + //! Copies data from a batch of given source buffers to their corresponding destination buffer. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! .. note:: + //! + //! If any input buffer aliases memory from any output buffer the behavior is undefined. + //! If any output buffer aliases memory of another output buffer the behavior is undefined. + //! Input buffers can alias one another. + //! + //! Snippet + //! +++++++ + //! + //! The code snippet below illustrates usage of DeviceMemcpy::Batched with an environment: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_memcpy_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin memcpy-batched-env + //! :end-before: example-end memcpy-batched-env + //! + //! @endrst + //! + //! @tparam InputBufferIt + //! **[inferred]** Device-accessible random-access input iterator type providing the pointers to + //! the source memory buffers + //! + //! @tparam OutputBufferIt + //! **[inferred]** Device-accessible random-access input iterator type providing the pointers to + //! the destination memory buffers + //! + //! @tparam BufferSizeIteratorT + //! **[inferred]** Device-accessible random-access input iterator type providing the number of bytes + //! to be copied for each pair of buffers + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] input_buffer_it + //! Device-accessible iterator providing the pointers to the source memory buffers + //! + //! @param[in] output_buffer_it + //! Device-accessible iterator providing the pointers to the destination memory buffers + //! + //! @param[in] buffer_sizes + //! Device-accessible iterator providing the number of bytes to be copied for each pair of buffers + //! + //! @param[in] num_buffers + //! The total number of buffer pairs + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + Batched(InputBufferIt input_buffer_it, + OutputBufferIt output_buffer_it, + BufferSizeIteratorT buffer_sizes, + ::cuda::std::int64_t num_buffers, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceMemcpy::Batched"); + static_assert(::cuda::std::is_pointer_v>, + "DeviceMemcpy::Batched only supports copying of memory buffers." + "Please consider using DeviceCopy::Batched instead."); + static_assert(::cuda::std::is_pointer_v>, + "DeviceMemcpy::Batched only supports copying of memory buffers." + "Please consider using DeviceCopy::Batched instead."); + + using BlockOffsetT = uint32_t; + using default_policy_selector = detail::batch_memcpy::policy_selector; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::batch_memcpy::dispatch( + storage, bytes, input_buffer_it, output_buffer_it, buffer_sizes, num_buffers, stream, policy_selector); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_merge.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_merge.cuh new file mode 100644 index 00000000..6803d000 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_merge.cuh @@ -0,0 +1,471 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceMerge provides device-wide, parallel operations for merging two sorted sequences of values (called keys) or +//! key-value pairs in device-accessible memory. The sorting order is determined by a comparison functor (default: +//! less-than), which has to establish a `strict weak ordering +//! `_. +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceMerge that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::MergePolicy`, as shown in the +//! example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin merge-keys-policy-selector +//! :end-before: example-end merge-keys-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin merge-keys-tuning +//! :end-before: example-end merge-keys-tuning +//! +//! @endrst +struct DeviceMerge +{ + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Merges two sorted sequences of values (called keys) into a sorted output sequence. Merging is unstable, + //! which means any two equivalent values (neither value is ordered before the other) may be written to the output + //! sequence in any order. + //! + //! .. versionadded:: 2.7.0 + //! First appears in CUDA Toolkit 12.8. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! The code snippet below illustrates the merging of two device vectors of `int` keys. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin merge-keys + //! :end-before: example-end merge-keys + //! + //! @endrst + //! + //! @tparam KeyIteratorIn1 **[deduced]** Random access iterator to the first sorted input sequence. Must have the same + //! value type as KeyIteratorIn2. + //! @tparam KeyIteratorIn2 **[deduced]** Random access iterator to the second sorted input sequence. Must have the + //! same value type as KeyIteratorIn1. + //! @tparam KeyIteratorOut **[deduced]** Random access iterator to the output sequence. + //! @tparam CompareOp **[deduced]** Binary predicate to compare the input iterator's value types. Must have a + //! signature equivalent to `bool operator()(Key lhs, Key rhs)` and establish a [strict weak ordering]. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation. + //! @param[in] keys_in1 Iterator to the beginning of the first sorted input sequence. + //! @param[in] num_keys1 Number of keys in the first input sequence. + //! @param[in] keys_in2 Iterator to the beginning of the second sorted input sequence. + //! @param[in] num_keys2 Number of keys in the second input sequence. + //! @param[out] keys_out Iterator to the beginning of the output sequence. + //! @param[in] compare_op Comparison function object, returning true if the first argument is ordered before the + //! second. Must establish a [strict weak ordering]. + //! @param[in] stream **[optional]** CUDA stream to launch kernels into. Default is stream0. + //! + //! [strict weak ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + template > + CUB_RUNTIME_FUNCTION static cudaError_t MergeKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIteratorIn1 keys_in1, + ::cuda::std::int64_t num_keys1, + KeyIteratorIn2 keys_in2, + ::cuda::std::int64_t num_keys2, + KeyIteratorOut keys_out, + CompareOp compare_op = {}, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceMerge::MergeKeys"); + // offset type is just int64_t + return detail::merge::dispatch( + d_temp_storage, + temp_storage_bytes, + keys_in1, + static_cast(nullptr), + num_keys1, + keys_in2, + static_cast(nullptr), + num_keys2, + keys_out, + static_cast(nullptr), + compare_op, + stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Merges two sorted sequences of values (called keys) into a sorted output sequence. Merging is unstable, + //! which means any two equivalent values (neither value is ordered before the other) may be written to the output + //! sequence in any order. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! Snippet + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin merge-keys-env + //! :end-before: example-end merge-keys-env + //! + //! @endrst + //! + //! @tparam KeyIteratorIn1 + //! **[deduced]** Random access iterator to the first sorted input sequence. Must have the same + //! value type as KeyIteratorIn2. + //! + //! @tparam KeyIteratorIn2 + //! **[deduced]** Random access iterator to the second sorted input sequence. Must have the + //! same value type as KeyIteratorIn1. + //! + //! @tparam KeyIteratorOut + //! **[deduced]** Random access iterator to the output sequence. + //! + //! @tparam CompareOp + //! **[deduced]** Binary predicate to compare the input iterator's value types. Must have a + //! signature equivalent to `bool operator()(Key lhs, Key rhs)` and establish a [strict weak ordering]. + //! + //! @tparam EnvT + //! **[deduced]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] keys_in1 + //! Iterator to the beginning of the first sorted input sequence. + //! + //! @param[in] num_keys1 + //! Number of keys in the first input sequence. + //! + //! @param[in] keys_in2 + //! Iterator to the beginning of the second sorted input sequence. + //! + //! @param[in] num_keys2 + //! Number of keys in the second input sequence. + //! + //! @param[out] keys_out + //! Iterator to the beginning of the output sequence. + //! + //! @param[in] compare_op + //! Comparison function object, returning true if the first argument is ordered before the + //! second. Must establish a [strict weak ordering]. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! [strict weak ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + template < + typename KeyIteratorIn1, + typename KeyIteratorIn2, + typename KeyIteratorOut, + typename CompareOp = ::cuda::std::less<>, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t< + !::cuda::std::is_same_v && !::cuda::std::is_same_v, + int> = 0, + ::cuda::std::enable_if_t<::cuda::std::indirect_binary_predicate, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MergeKeys( + KeyIteratorIn1 keys_in1, + ::cuda::std::int64_t num_keys1, + KeyIteratorIn2 keys_in2, + ::cuda::std::int64_t num_keys2, + KeyIteratorOut keys_out, + CompareOp compare_op = {}, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceMerge::MergeKeys"); + + using default_policy_selector = + detail::merge::policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return detail::merge::dispatch( + d_temp_storage, + temp_storage_bytes, + keys_in1, + static_cast(nullptr), + num_keys1, + keys_in2, + static_cast(nullptr), + num_keys2, + keys_out, + static_cast(nullptr), + compare_op, + stream, + policy_selector); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Merges two sorted sequences of key-value pairs into a sorted output sequence. Merging is unstable, + //! which means any two equivalent values (neither value is ordered before the other) may be written to the output + //! sequence in any order. + //! + //! .. versionadded:: 2.7.0 + //! First appears in CUDA Toolkit 12.8. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! The code snippet below illustrates the merging of two device vectors of `int` keys. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin merge-pairs + //! :end-before: example-end merge-pairs + //! + //! @endrst + //! + //! @tparam KeyIteratorIn1 **[deduced]** Random access iterator to the keys of the first sorted input sequence. Must + //! have the same value type as KeyIteratorIn2. + //! @tparam ValueIteratorIn1 **[deduced]** Random access iterator to the values of the first sorted input sequence. + //! Must have the same value type as ValueIteratorIn2. + //! @tparam KeyIteratorIn2 **[deduced]** Random access iterator to the second sorted input sequence. Must have the + //! same value type as KeyIteratorIn1. + //! @tparam ValueIteratorIn2 **[deduced]** Random access iterator to the values of the second sorted input sequence. + //! Must have the same value type as ValueIteratorIn1. + //! @tparam KeyIteratorOut **[deduced]** Random access iterator to the keys of the output sequence. + //! @tparam ValueIteratorOut **[deduced]** Random access iterator to the values of the output sequence. + //! @tparam CompareOp **[deduced]** Binary predicate to compare the key input iterator's value types. Must have a + //! signature equivalent to `bool operator()(Key lhs, Key rhs)` and establish a [strict weak ordering]. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation. + //! @param[in] keys_in1 Iterator to the beginning of the keys of the first sorted input sequence. + //! @param[in] values_in1 Iterator to the beginning of the values of the first sorted input sequence. + //! @param[in] num_pairs1 Number of key-value pairs in the first input sequence. + //! @param[in] keys_in2 Iterator to the beginning of the keys of the second sorted input sequence. + //! @param[in] values_in2 Iterator to the beginning of the values of the second sorted input sequence. + //! @param[in] num_pairs2 Number of key-value pairs in the second input sequence. + //! @param[out] keys_out Iterator to the beginning of the keys of the output sequence. + //! @param[out] values_out Iterator to the beginning of the values of the output sequence. + //! @param[in] compare_op Comparison function object, returning true if the first argument is ordered before the + //! second. Must establish a [strict weak ordering]. + //! @param[in] stream **[optional]** CUDA stream to launch kernels into. Default is stream0. + //! + //! [strict weak ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + template > + CUB_RUNTIME_FUNCTION static cudaError_t MergePairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIteratorIn1 keys_in1, + ValueIteratorIn1 values_in1, + ::cuda::std::int64_t num_pairs1, + KeyIteratorIn2 keys_in2, + ValueIteratorIn2 values_in2, + ::cuda::std::int64_t num_pairs2, + KeyIteratorOut keys_out, + ValueIteratorOut values_out, + CompareOp compare_op = {}, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceMerge::MergePairs"); + // offset type is just int64_t + return detail::merge::dispatch( + d_temp_storage, + temp_storage_bytes, + keys_in1, + values_in1, + num_pairs1, + keys_in2, + values_in2, + num_pairs2, + keys_out, + values_out, + compare_op, + stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Merges two sorted sequences of key-value pairs into a sorted output sequence. Merging is unstable, + //! which means any two equivalent values (neither value is ordered before the other) may be written to the output + //! sequence in any order. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin merge-pairs-env + //! :end-before: example-end merge-pairs-env + //! + //! @endrst + //! + //! @tparam KeyIteratorIn1 + //! **[deduced]** Random access iterator to the keys of the first sorted input sequence. Must + //! have the same value type as KeyIteratorIn2. + //! + //! @tparam ValueIteratorIn1 + //! **[deduced]** Random access iterator to the values of the first sorted input sequence. + //! Must have the same value type as ValueIteratorIn2. + //! + //! @tparam KeyIteratorIn2 + //! **[deduced]** Random access iterator to the second sorted input sequence. Must have the + //! same value type as KeyIteratorIn1. + //! + //! @tparam ValueIteratorIn2 + //! **[deduced]** Random access iterator to the values of the second sorted input sequence. + //! Must have the same value type as ValueIteratorIn1. + //! + //! @tparam KeyIteratorOut + //! **[deduced]** Random access iterator to the keys of the output sequence. + //! + //! @tparam ValueIteratorOut + //! **[deduced]** Random access iterator to the values of the output sequence. + //! + //! @tparam CompareOp + //! **[deduced]** Binary predicate to compare the key input iterator's value types. Must have a + //! signature equivalent to `bool operator()(Key lhs, Key rhs)` and establish a [strict weak ordering]. + //! + //! @tparam EnvT + //! **[deduced]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] keys_in1 + //! Iterator to the beginning of the keys of the first sorted input sequence. + //! + //! @param[in] values_in1 + //! Iterator to the beginning of the values of the first sorted input sequence. + //! + //! @param[in] num_pairs1 + //! Number of key-value pairs in the first input sequence. + //! + //! @param[in] keys_in2 + //! Iterator to the beginning of the keys of the second sorted input sequence. + //! + //! @param[in] values_in2 + //! Iterator to the beginning of the values of the second sorted input sequence. + //! + //! @param[in] num_pairs2 + //! Number of key-value pairs in the second input sequence. + //! + //! @param[out] keys_out + //! Iterator to the beginning of the keys of the output sequence. + //! + //! @param[out] values_out + //! Iterator to the beginning of the values of the output sequence. + //! + //! @param[in] compare_op + //! Comparison function object, returning true if the first argument is ordered before the + //! second. Must establish a [strict weak ordering]. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + //! [strict weak ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + template < + typename KeyIteratorIn1, + typename ValueIteratorIn1, + typename KeyIteratorIn2, + typename ValueIteratorIn2, + typename KeyIteratorOut, + typename ValueIteratorOut, + typename CompareOp = ::cuda::std::less<>, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t< + !::cuda::std::is_same_v && !::cuda::std::is_same_v, + int> = 0, + ::cuda::std::enable_if_t<::cuda::std::indirect_binary_predicate, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t MergePairs( + KeyIteratorIn1 keys_in1, + ValueIteratorIn1 values_in1, + ::cuda::std::int64_t num_pairs1, + KeyIteratorIn2 keys_in2, + ValueIteratorIn2 values_in2, + ::cuda::std::int64_t num_pairs2, + KeyIteratorOut keys_out, + ValueIteratorOut values_out, + CompareOp compare_op = {}, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceMerge::MergePairs"); + using default_policy_selector = detail::merge:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return detail::merge::dispatch( + d_temp_storage, + temp_storage_bytes, + keys_in1, + values_in1, + num_pairs1, + keys_in2, + values_in2, + num_pairs2, + keys_out, + values_out, + compare_op, + stream, + policy_selector); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_merge_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_merge_sort.cuh new file mode 100644 index 00000000..ecc5572f --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_merge_sort.cuh @@ -0,0 +1,1607 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +/** + * @brief DeviceMergeSort provides device-wide, parallel operations for + * computing a merge sort across a sequence of data items residing within + * device-accessible memory. + * + * @par Overview + * - DeviceMergeSort arranges items into ascending order using a comparison + * functor with less-than semantics. Merge sort can handle arbitrary types (as + * long as a value of these types is a model of [LessThan Comparable]) and + * comparison functors, but is slower than DeviceRadixSort when sorting + * arithmetic types into ascending/descending order. + * - Another difference from RadixSort is the fact that DeviceMergeSort can + * handle arbitrary random-access iterators, as shown below. + * + * @par A Simple Example + * @par + * The code snippet below illustrates a thrust reverse iterator usage. + * @par + * @code + * #include // or equivalently + * + * struct CustomLess + * { + * template + * __device__ bool operator()(const DataType &lhs, const DataType &rhs) + * { + * return lhs < rhs; + * } + * }; + * + * // Declare, allocate, and initialize device-accessible pointers + * // for sorting data + * thrust::device_vector d_keys(num_items); + * thrust::device_vector d_values(num_items); + * // ... + * + * // Initialize iterator + * using KeyIterator = typename thrust::device_vector::iterator; + * cuda::std::reverse_iterator reverse_iter(d_keys.end()); + * + * // Determine temporary device storage requirements + * size_t temp_storage_bytes = 0; + * cub::DeviceMergeSort::SortPairs( + * nullptr, + * temp_storage_bytes, + * reverse_iter, + * thrust::raw_pointer_cast(d_values.data()), + * num_items, + * CustomLess()); + * + * // Allocate temporary storage + * cudaMalloc(&d_temp_storage, temp_storage_bytes); + * + * // Run sorting operation + * cub::DeviceMergeSort::SortPairs( + * d_temp_storage, + * temp_storage_bytes, + * reverse_iter, + * thrust::raw_pointer_cast(d_values.data()), + * num_items, + * CustomLess()); + * @endcode + * + * @rst + * + * Tuning + * +++++++++++++++++++++++++++++++++++++++++++++ + * + * All algorithms in DeviceMergeSort that accept an environment can be tuned by passing a custom :ref:`policy selector + * ` that returns a :cpp:struct:`cub::MergeSortPolicy`, as shown in the example below: + * + * .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + * :language: c++ + * :dedent: + * :start-after: example-begin sort-pairs-policy-selector + * :end-before: example-end sort-pairs-policy-selector + * + * .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + * :language: c++ + * :dedent: + * :start-after: example-begin sort-pairs-tuning + * :end-before: example-end sort-pairs-tuning + * @endrst + * + * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable + */ +struct DeviceMergeSort +{ +private: + // Name reported for NVTX ranges + _CCCL_HOST_DEVICE static constexpr auto GetName() -> const char* + { + return "cub::DeviceMergeSort"; + } + + // Internal version without NVTX range + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsNoNVTX( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIteratorT d_keys, + ValueIteratorT d_values, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + using ChooseOffsetT = detail::choose_offset_t; + + return detail::merge_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + d_keys, + d_values, + static_cast(num_items), + compare_op, + stream); + } + +public: + /** + * @brief Sorts items using a merge sorting method. + * + * @par + * SortPairs is not guaranteed to be stable. That is, suppose that i and j are + * equivalent: neither one is less than the other. It is not guaranteed + * that the relative order of these two elements will be preserved by sort. + * + * @par Snippet + * The code snippet below illustrates the sorting of a device vector of `int` + * keys with associated vector of `int` values. + * @par + * @code + * #include + * // or equivalently + * + * // Declare, allocate, and initialize device-accessible pointers for + * // sorting data + * int num_items; // e.g., 7 + * int *d_keys; // e.g., [8, 6, 6, 5, 3, 0, 9] + * int *d_values; // e.g., [0, 1, 2, 3, 4, 5, 6] + * ... + * + * // Initialize comparator + * CustomOpT custom_op; + * + * // Determine temporary device storage requirements + * void *d_temp_storage = nullptr; + * size_t temp_storage_bytes = 0; + * cub::DeviceMergeSort::SortPairs( + * d_temp_storage, temp_storage_bytes, + * d_keys, d_values, num_items, custom_op); + * + * // Allocate temporary storage + * cudaMalloc(&d_temp_storage, temp_storage_bytes); + * + * // Run sorting operation + * cub::DeviceMergeSort::SortPairs( + * d_temp_storage, temp_storage_bytes, + * d_keys, d_values, num_items, custom_op); + * + * // d_keys <-- [0, 3, 5, 6, 6, 8, 9] + * // d_values <-- [5, 4, 3, 2, 1, 0, 6] + * + * @endcode + * + * @tparam KeyIteratorT + * is a model of [Random Access Iterator]. `KeyIteratorT` is mutable, and + * its `value_type` is a model of [LessThan Comparable]. This `value_type`'s + * ordering relation is a *strict weak ordering* as defined in + * the [LessThan Comparable] requirements. + * + * @tparam ValueIteratorT + * is a model of [Random Access Iterator], and `ValueIteratorT` is mutable. + * + * @tparam OffsetT + * is an integer type for global offsets. + * + * @tparam CompareOpT + * is a type of callable object with the signature + * `bool operator()(KeyT lhs, KeyT rhs)` that models + * the [Strict Weak Ordering] concept. + * + * @param[in] d_temp_storage + * @devicestorage + * + * @param[in,out] temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param[in,out] d_keys + * Pointer to the input sequence of unsorted input keys + * + * @param[in,out] d_values + * Pointer to the input sequence of unsorted input values + * + * @param[in] num_items + * Number of items to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. Default is + * stream0. + * + * [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIteratorT d_keys, + ValueIteratorT d_values, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return SortPairsNoNVTX(d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, compare_op, stream); + } + + //! @rst + //! Sorts items using a merge sorting method. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - SortPairs is not guaranteed to be stable. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-pairs-env + //! :end-before: example-end sort-pairs-env + //! + //! @endrst + //! + //! @tparam KeyIteratorT + //! **[inferred]** Random-access iterator type for keys @iterator + //! + //! @tparam ValueIteratorT + //! **[inferred]** Random-access iterator type for values @iterator + //! + //! @tparam OffsetT + //! **[inferred]** Integer type for offsets + //! + //! @tparam CompareOpT + //! **[inferred]** Comparison function object type + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type + //! + //! @param[in,out] d_keys + //! Keys to sort + //! + //! @param[in,out] d_values + //! Values corresponding to keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] compare_op + //! Comparison function object + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( + KeyIteratorT d_keys, ValueIteratorT d_values, OffsetT num_items, CompareOpT compare_op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using ChooseOffsetT = detail::choose_offset_t; + using default_policy_selector = detail::merge_sort::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::merge_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + d_keys, + d_values, + static_cast(num_items), + compare_op, + stream, + policy_selector); + }); + } + + /** + * @brief Sorts items using a merge sorting method. + * + * @par + * - SortPairsCopy is not guaranteed to be stable. That is, suppose + * that `i` and `j` are equivalent: neither one is less than the + * other. It is not guaranteed that the relative order of these + * two elements will be preserved by sort. + * - Input arrays `d_input_keys` and `d_input_values` are not modified. + * - Note that the behavior is undefined if the input and output ranges + * overlap in any way. + * + * @par Snippet + * The code snippet below illustrates the sorting of a device vector of + * `int` keys with associated vector of `int` values. + * @par + * @code + * #include + * // or equivalently + * + * // Declare, allocate, and initialize device-accessible pointers + * // for sorting data + * int num_items; // e.g., 7 + * int *d_keys; // e.g., [8, 6, 6, 5, 3, 0, 9] + * int *d_values; // e.g., [0, 1, 2, 3, 4, 5, 6] + * ... + * + * // Initialize comparator + * CustomOpT custom_op; + * + * // Determine temporary device storage requirements + * void *d_temp_storage = nullptr; + * size_t temp_storage_bytes = 0; + * cub::DeviceMergeSort::SortPairsCopy( + * d_temp_storage, temp_storage_bytes, + * d_keys, d_values, num_items, custom_op); + * + * // Allocate temporary storage + * cudaMalloc(&d_temp_storage, temp_storage_bytes); + * + * // Run sorting operation + * cub::DeviceMergeSort::SortPairsCopy( + * d_temp_storage, temp_storage_bytes, + * d_keys, d_values, num_items, custom_op); + * + * // d_keys <-- [0, 3, 5, 6, 6, 8, 9] + * // d_values <-- [5, 4, 3, 2, 1, 0, 6] + * + * @endcode + * + * @tparam KeyInputIteratorT + * is a model of [Random Access Iterator]. Its `value_type` is a model of + * [LessThan Comparable]. This `value_type`'s ordering relation is a + * *strict weak ordering* as defined in the [LessThan Comparable] + * requirements. + * + * @tparam ValueInputIteratorT + * is a model of [Random Access Iterator]. + * + * @tparam KeyIteratorT + * is a model of [Random Access Iterator]. `KeyIteratorT` is mutable, and + * its `value_type` is a model of [LessThan Comparable]. This `value_type`'s + * ordering relation is a *strict weak ordering* as defined in + * the [LessThan Comparable] requirements. + * + * @tparam ValueIteratorT + * is a model of [Random Access Iterator], and `ValueIteratorT` is mutable. + * + * @tparam OffsetT + * is an integer type for global offsets. + * + * @tparam CompareOpT + * is a type of callable object with the signature + * `bool operator()(KeyT lhs, KeyT rhs)` that models + * the [Strict Weak Ordering] concept. + * + * @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_input_keys + * Pointer to the input sequence of unsorted input keys + * + * @param[in] d_input_values + * Pointer to the input sequence of unsorted input values + * + * @param[out] d_output_keys + * Pointer to the output sequence of sorted input keys + * + * @param[out] d_output_values + * Pointer to the output sequence of sorted input values + * + * @param[in] num_items + * Number of items to sort + * + * @param[in] compare_op + * Comparison function object which returns `true` if the first argument is + * ordered before the second + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. Default is + * stream0. + * + * [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsCopy( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_input_keys, + ValueInputIteratorT d_input_values, + KeyIteratorT d_output_keys, + ValueIteratorT d_output_values, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + using ChooseOffsetT = detail::choose_offset_t; + + return detail::merge_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_input_keys, + d_input_values, + d_output_keys, + d_output_values, + static_cast(num_items), + compare_op, + stream); + } + + //! @rst + //! Sorts items using a merge sorting method. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - SortPairsCopy is not guaranteed to be stable. + //! - Input arrays ``d_input_keys`` and ``d_input_values`` are not modified. + //! - The behavior is undefined if the input and output ranges overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-pairs-copy-env + //! :end-before: example-end sort-pairs-copy-env + //! + //! @endrst + //! + //! @tparam KeyInputIteratorT + //! **[inferred]** Random-access iterator type for input keys @iterator + //! + //! @tparam ValueInputIteratorT + //! **[inferred]** Random-access iterator type for input values @iterator + //! + //! @tparam KeyIteratorT + //! **[inferred]** Random-access iterator type for output keys @iterator + //! + //! @tparam ValueIteratorT + //! **[inferred]** Random-access iterator type for output values @iterator + //! + //! @tparam OffsetT + //! **[inferred]** Integer type for offsets + //! + //! @tparam CompareOpT + //! **[inferred]** Comparison function object type + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type + //! + //! @param[in] d_input_keys + //! Pointer to the input sequence of unsorted input keys + //! + //! @param[in] d_input_values + //! Pointer to the input sequence of unsorted input values + //! + //! @param[out] d_output_keys + //! Pointer to the output sequence of sorted input keys + //! + //! @param[out] d_output_values + //! Pointer to the output sequence of sorted input values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] compare_op + //! Comparison function object + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsCopy( + KeyInputIteratorT d_input_keys, + ValueInputIteratorT d_input_values, + KeyIteratorT d_output_keys, + ValueIteratorT d_output_values, + OffsetT num_items, + CompareOpT compare_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using ChooseOffsetT = detail::choose_offset_t; + using default_policy_selector = detail::merge_sort::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::merge_sort::dispatch( + storage, + bytes, + d_input_keys, + d_input_values, + d_output_keys, + d_output_values, + static_cast(num_items), + compare_op, + stream, + policy_selector); + }); + } + +private: + // Internal version without NVTX range + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysNoNVTX( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIteratorT d_keys, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + using ChooseOffsetT = detail::choose_offset_t; + + return detail::merge_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + static_cast(nullptr), + d_keys, + static_cast(nullptr), + static_cast(num_items), + compare_op, + stream); + } + +public: + /** + * @brief Sorts items using a merge sorting method. + * + * @par + * SortKeys is not guaranteed to be stable. That is, suppose that `i` and `j` + * are equivalent: neither one is less than the other. It is not guaranteed + * that the relative order of these two elements will be preserved by sort. + * + * @par Snippet + * The code snippet below illustrates the sorting of a device vector of `int` + * keys. + * @par + * @code + * #include + * // or equivalently + * + * // Declare, allocate, and initialize device-accessible pointers + * // for sorting data + * int num_items; // e.g., 7 + * int *d_keys; // e.g., [8, 6, 7, 5, 3, 0, 9] + * ... + * + * // Initialize comparator + * CustomOpT custom_op; + * + * // Determine temporary device storage requirements + * void *d_temp_storage = nullptr; + * size_t temp_storage_bytes = 0; + * cub::DeviceMergeSort::SortKeys( + * d_temp_storage, temp_storage_bytes, + * d_keys, num_items, custom_op); + * + * // Allocate temporary storage + * cudaMalloc(&d_temp_storage, temp_storage_bytes); + * + * // Run sorting operation + * cub::DeviceMergeSort::SortKeys( + * d_temp_storage, temp_storage_bytes, + * d_keys, num_items, custom_op); + * + * // d_keys <-- [0, 3, 5, 6, 7, 8, 9] + * @endcode + * + * @tparam KeyIteratorT + * is a model of [Random Access Iterator]. `KeyIteratorT` is mutable, and + * its `value_type` is a model of [LessThan Comparable]. This `value_type`'s + * ordering relation is a *strict weak ordering* as defined in + * the [LessThan Comparable] requirements. + * + * @tparam OffsetT + * is an integer type for global offsets. + * + * @tparam CompareOpT + * is a type of callable object with the signature + * `bool operator()(KeyT lhs, KeyT rhs)` that models + * the [Strict Weak Ordering] concept. + * + * @param[in] d_temp_storage + * @devicestorage + * + * @param[in,out] temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param[in,out] d_keys + * Pointer to the input sequence of unsorted input keys + * + * @param[in] num_items + * Number of items to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. Default is + * stream0. + * + * [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIteratorT d_keys, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return SortKeysNoNVTX(d_temp_storage, temp_storage_bytes, d_keys, num_items, compare_op, stream); + } + + //! @rst + //! Sorts keys using a merge sorting method. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - SortKeys is not guaranteed to be stable. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-keys-env + //! :end-before: example-end sort-keys-env + //! + //! @endrst + //! + //! @tparam KeyIteratorT + //! **[inferred]** Random-access iterator type for keys @iterator + //! + //! @tparam OffsetT + //! **[inferred]** Integer type for offsets + //! + //! @tparam CompareOpT + //! **[inferred]** Comparison function object type + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type + //! + //! @param[in,out] d_keys + //! Keys to sort + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] compare_op + //! Comparison function object + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + SortKeys(KeyIteratorT d_keys, OffsetT num_items, CompareOpT compare_op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using ChooseOffsetT = detail::choose_offset_t; + using default_policy_selector = detail::merge_sort::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::merge_sort::dispatch( + storage, + bytes, + d_keys, + static_cast(nullptr), + d_keys, + static_cast(nullptr), + static_cast(num_items), + compare_op, + stream, + policy_selector); + }); + } + +private: + // Internal version without NVTX range + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysCopyNoNVTX( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_input_keys, + KeyIteratorT d_output_keys, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + using ChooseOffsetT = detail::choose_offset_t; + + return detail::merge_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_input_keys, + static_cast(nullptr), + d_output_keys, + static_cast(nullptr), + static_cast(num_items), + compare_op, + stream); + } + +public: + /** + * @brief Sorts items using a merge sorting method. + * + * @par + * - SortKeysCopy is not guaranteed to be stable. That is, suppose that `i` + * and `j` are equivalent: neither one is less than the other. It is not + * guaranteed that the relative order of these two elements will be + * preserved by sort. + * - Input array d_input_keys is not modified. + * - Note that the behavior is undefined if the input and output ranges + * overlap in any way. + * + * @par Snippet + * The code snippet below illustrates the sorting of a device vector of + * `int` keys. + * @par + * @code + * #include + * // or equivalently + * + * // Declare, allocate, and initialize device-accessible pointers for + * // sorting data + * int num_items; // e.g., 7 + * int *d_input_keys; // e.g., [8, 6, 7, 5, 3, 0, 9] + * int *d_output_keys; // must hold at least num_items elements + * ... + * + * // Initialize comparator + * CustomOpT custom_op; + * + * // Determine temporary device storage requirements + * void *d_temp_storage = nullptr; + * size_t temp_storage_bytes = 0; + * cub::DeviceMergeSort::SortKeysCopy( + * d_temp_storage, temp_storage_bytes, + * d_input_keys, d_output_keys, num_items, custom_op); + * + * // Allocate temporary storage + * cudaMalloc(&d_temp_storage, temp_storage_bytes); + * + * // Run sorting operation + * cub::DeviceMergeSort::SortKeysCopy( + * d_temp_storage, temp_storage_bytes, + * d_input_keys, d_output_keys, num_items, custom_op); + * + * // d_output_keys <-- [0, 3, 5, 6, 7, 8, 9] + * @endcode + * + * @tparam KeyInputIteratorT + * is a model of [Random Access Iterator]. Its `value_type` is a model of + * [LessThan Comparable]. This `value_type`'s ordering relation is a + * *strict weak ordering* as defined in the [LessThan Comparable] + * requirements. + * + * @tparam KeyIteratorT + * is a model of [Random Access Iterator]. `KeyIteratorT` is mutable, and + * its `value_type` is a model of [LessThan Comparable]. This `value_type`'s + * ordering relation is a *strict weak ordering* as defined in + * the [LessThan Comparable] requirements. + * + * @tparam OffsetT + * is an integer type for global offsets. + * + * @tparam CompareOpT + * is a type of callable object with the signature + * `bool operator()(KeyT lhs, KeyT rhs)` that models + * the [Strict Weak Ordering] concept. + * + * @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_input_keys + * Pointer to the input sequence of unsorted input keys + * + * @param[out] d_output_keys + * Pointer to the output sequence of sorted input keys + * + * @param[in] num_items + * Number of items to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. Default is + * stream0. + * + * [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysCopy( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_input_keys, + KeyIteratorT d_output_keys, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return SortKeysCopyNoNVTX( + d_temp_storage, temp_storage_bytes, d_input_keys, d_output_keys, num_items, compare_op, stream); + } + + //! @rst + //! Sorts keys using a merge sorting method. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - SortKeysCopy is not guaranteed to be stable. + //! - Input array ``d_input_keys`` is not modified. + //! - The behavior is undefined if the input and output ranges overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-keys-copy-env + //! :end-before: example-end sort-keys-copy-env + //! + //! @endrst + //! + //! @tparam KeyInputIteratorT + //! **[inferred]** Random-access iterator type for input keys @iterator + //! + //! @tparam KeyIteratorT + //! **[inferred]** Random-access iterator type for output keys @iterator + //! + //! @tparam OffsetT + //! **[inferred]** Integer type for offsets + //! + //! @tparam CompareOpT + //! **[inferred]** Comparison function object type + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type + //! + //! @param[in] d_input_keys + //! Pointer to the input sequence of unsorted input keys + //! + //! @param[out] d_output_keys + //! Pointer to the output sequence of sorted input keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] compare_op + //! Comparison function object + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysCopy( + KeyInputIteratorT d_input_keys, + KeyIteratorT d_output_keys, + OffsetT num_items, + CompareOpT compare_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using ChooseOffsetT = detail::choose_offset_t; + using default_policy_selector = detail::merge_sort::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::merge_sort::dispatch( + storage, + bytes, + d_input_keys, + static_cast(nullptr), + d_output_keys, + static_cast(nullptr), + static_cast(num_items), + compare_op, + stream, + policy_selector); + }); + } + + /** + * @brief Sorts items using a merge sorting method. + * + * @par + * StableSortPairs is stable: it preserves the relative ordering of equivalent + * elements. That is, if x and y are elements such that x precedes y, + * and if the two elements are equivalent (neither x < y nor y < x) then + * a postcondition of stable_sort is that x still precedes y. + * + * @par Snippet + * The code snippet below illustrates the sorting of a device vector of `int` + * keys with associated vector of `int` values. + * @par + * @code + * #include + * // or equivalently + * + * // Declare, allocate, and initialize device-accessible pointers for + * // sorting data + * int num_items; // e.g., 7 + * int *d_keys; // e.g., [8, 6, 6, 5, 3, 0, 9] + * int *d_values; // e.g., [0, 1, 2, 3, 4, 5, 6] + * ... + * + * // Initialize comparator + * CustomOpT custom_op; + * + * // Determine temporary device storage requirements + * void *d_temp_storage = nullptr; + * size_t temp_storage_bytes = 0; + * cub::DeviceMergeSort::StableSortPairs( + * d_temp_storage, temp_storage_bytes, + * d_keys, d_values, num_items, custom_op); + * + * // Allocate temporary storage + * cudaMalloc(&d_temp_storage, temp_storage_bytes); + * + * // Run sorting operation + * cub::DeviceMergeSort::StableSortPairs( + * d_temp_storage, temp_storage_bytes, + * d_keys, d_values, num_items, custom_op); + * + * // d_keys <-- [0, 3, 5, 6, 6, 8, 9] + * // d_values <-- [5, 4, 3, 1, 2, 0, 6] + * @endcode + * + * @tparam KeyIteratorT + * is a model of [Random Access Iterator]. `KeyIteratorT` is mutable, and + * its `value_type` is a model of [LessThan Comparable]. This `value_type`'s + * ordering relation is a *strict weak ordering* as defined in + * the [LessThan Comparable] requirements. + * + * @tparam ValueIteratorT + * is a model of [Random Access Iterator], and `ValueIteratorT` is mutable. + * + * @tparam OffsetT + * is an integer type for global offsets. + * + * @tparam CompareOpT + * is a type of callable object with the signature + * `bool operator()(KeyT lhs, KeyT rhs)` that models + * the [Strict Weak Ordering] concept. + * + * @param[in] d_temp_storage + * @devicestorage + * + * @param[in,out] temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param[in,out] d_keys + * Pointer to the input sequence of unsorted input keys + * + * @param[in,out] d_values + * Pointer to the input sequence of unsorted input values + * + * @param[in] num_items + * Number of items to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. Default is + * stream0. + * + * [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIteratorT d_keys, + ValueIteratorT d_values, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + return SortPairsNoNVTX( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, compare_op, stream); + } + + //! @rst + //! Stably sorts items using a merge sorting method. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - StableSortPairs preserves the relative ordering of equivalent elements. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-pairs-env + //! :end-before: example-end stable-sort-pairs-env + //! + //! @endrst + //! + //! @tparam KeyIteratorT + //! **[inferred]** Random-access iterator type for keys @iterator + //! + //! @tparam ValueIteratorT + //! **[inferred]** Random-access iterator type for values @iterator + //! + //! @tparam OffsetT + //! **[inferred]** Integer type for offsets + //! + //! @tparam CompareOpT + //! **[inferred]** Comparison function object type + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type + //! + //! @param[in,out] d_keys + //! Keys to sort + //! + //! @param[in,out] d_values + //! Values corresponding to keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] compare_op + //! Comparison function object + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t StableSortPairs( + KeyIteratorT d_keys, ValueIteratorT d_values, OffsetT num_items, CompareOpT compare_op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using ChooseOffsetT = detail::choose_offset_t; + using default_policy_selector = detail::merge_sort::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::merge_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + d_keys, + d_values, + static_cast(num_items), + compare_op, + stream, + policy_selector); + }); + } + + /** + * @brief Sorts items using a merge sorting method. + * + * @par + * StableSortKeys is stable: it preserves the relative ordering of equivalent + * elements. That is, if `x` and `y` are elements such that `x` precedes `y`, + * and if the two elements are equivalent (neither `x < y` nor `y < x`) then + * a postcondition of stable_sort is that `x` still precedes `y`. + * + * @par Snippet + * The code snippet below illustrates the sorting of a device vector of `int` + * keys. + * \par + * \code + * #include + * // or equivalently + * + * // Declare, allocate, and initialize device-accessible pointers for + * // sorting data + * int num_items; // e.g., 7 + * int *d_keys; // e.g., [8, 6, 7, 5, 3, 0, 9] + * ... + * + * // Initialize comparator + * CustomOpT custom_op; + * + * // Determine temporary device storage requirements + * void *d_temp_storage = nullptr; + * size_t temp_storage_bytes = 0; + * cub::DeviceMergeSort::StableSortKeys( + * d_temp_storage, temp_storage_bytes, + * d_keys, num_items, custom_op); + * + * // Allocate temporary storage + * cudaMalloc(&d_temp_storage, temp_storage_bytes); + * + * // Run sorting operation + * cub::DeviceMergeSort::StableSortKeys( + * d_temp_storage, temp_storage_bytes, + * d_keys, num_items, custom_op); + * + * // d_keys <-- [0, 3, 5, 6, 7, 8, 9] + * @endcode + * + * @tparam KeyIteratorT + * is a model of [Random Access Iterator]. `KeyIteratorT` is mutable, and + * its `value_type` is a model of [LessThan Comparable]. This `value_type`'s + * ordering relation is a *strict weak ordering* as defined in + * the [LessThan Comparable] requirements. + * + * @tparam OffsetT + * is an integer type for global offsets. + * + * @tparam CompareOpT + * is a type of callable object with the signature + * `bool operator()(KeyT lhs, KeyT rhs)` that models + * the [Strict Weak Ordering] concept. + * + * @param[in] d_temp_storage + * @devicestorage + * + * @param[in,out] temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param[in,out] d_keys + * Pointer to the input sequence of unsorted input keys + * + * @param[in] num_items + * Number of items to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. Default is + * stream0. + * + * [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIteratorT d_keys, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + return SortKeysNoNVTX( + d_temp_storage, temp_storage_bytes, d_keys, num_items, compare_op, stream); + } + + //! @rst + //! Stably sorts keys using a merge sorting method. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - StableSortKeys preserves the relative ordering of equivalent elements. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-keys-env + //! :end-before: example-end stable-sort-keys-env + //! + //! @endrst + //! + //! @tparam KeyIteratorT + //! **[inferred]** Random-access iterator type for keys @iterator + //! + //! @tparam OffsetT + //! **[inferred]** Integer type for offsets + //! + //! @tparam CompareOpT + //! **[inferred]** Comparison function object type + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type + //! + //! @param[in,out] d_keys + //! Keys to sort + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] compare_op + //! Comparison function object + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + StableSortKeys(KeyIteratorT d_keys, OffsetT num_items, CompareOpT compare_op, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using ChooseOffsetT = detail::choose_offset_t; + using default_policy_selector = detail::merge_sort::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::merge_sort::dispatch( + storage, + bytes, + d_keys, + static_cast(nullptr), + d_keys, + static_cast(nullptr), + static_cast(num_items), + compare_op, + stream, + policy_selector); + }); + } + + /** + * @brief Sorts items using a merge sorting method. + * + * @par + * - StableSortKeysCopy is stable: it preserves the relative ordering of equivalent + * elements. That is, if `x` and `y` are elements such that `x` precedes `y`, + * and if the two elements are equivalent (neither `x < y` nor `y < x`) then + * a postcondition of stable_sort is that `x` still precedes `y`. + * - Input array d_input_keys is not modified + * - Note that the behavior is undefined if the input and output ranges overlap + * in any way. + * + * @par Snippet + * The code snippet below illustrates the sorting of a device vector of `int` + * keys. + * \par + * \code + * #include + * // or equivalently + * + * // Declare, allocate, and initialize device-accessible pointers for + * // sorting data + * int num_items; // e.g., 7 + * int *d_input_keys; // e.g., [8, 6, 7, 5, 3, 0, 9] + * int *d_output_keys; // must hold at least num_items elements + * ... + * + * // Initialize comparator + * CustomOpT custom_op; + * + * // Determine temporary device storage requirements + * void *d_temp_storage = nullptr; + * size_t temp_storage_bytes = 0; + * cub::DeviceMergeSort::StableSortKeysCopy( + * d_temp_storage, temp_storage_bytes, + * d_input_keys, d_output_keys, num_items, custom_op); + * + * // Allocate temporary storage + * cudaMalloc(&d_temp_storage, temp_storage_bytes); + * + * // Run sorting operation + * cub::DeviceMergeSort::StableSortKeysCopy( + * d_temp_storage, temp_storage_bytes, + * d_input_keys, d_output_keys, num_items, custom_op); + * + * // d_output_keys <-- [0, 3, 5, 6, 7, 8, 9] + * @endcode + * + * @tparam KeyInputIteratorT + * is a model of [Random Access Iterator]. Its `value_type` is a model of + * [LessThan Comparable]. This `value_type`'s ordering relation is a + * *strict weak ordering* as defined in the [LessThan Comparable] + * requirements. + * + * @tparam KeyIteratorT + * is a model of [Random Access Iterator]. `KeyIteratorT` is mutable, and + * its `value_type` is a model of [LessThan Comparable]. This `value_type`'s + * ordering relation is a *strict weak ordering* as defined in + * the [LessThan Comparable] requirements. + * + * @tparam OffsetT + * is an integer type for global offsets. + * + * @tparam CompareOpT + * is a type of callable object with the signature + * `bool operator()(KeyT lhs, KeyT rhs)` that models + * the [Strict Weak Ordering] concept. + * + * @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_input_keys + * Pointer to the input sequence of unsorted input keys + * + * @param[out] d_output_keys + * Pointer to the output sequence of sorted input keys + * + * @param[in] num_items + * Number of elements in d_input_keys to sort + * + * @param[in] compare_op + * Comparison function object which returns true if the first argument is + * ordered before the second + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. Default is + * stream0. + * + * [Random Access Iterator]: https://en.cppreference.com/w/cpp/iterator/random_access_iterator + * [Strict Weak Ordering]: https://en.cppreference.com/w/cpp/concepts/strict_weak_order + * [LessThan Comparable]: https://en.cppreference.com/w/cpp/named_req/LessThanComparable + * + * @rst + * .. versionadded:: 2.2.0 + * First appears in CUDA Toolkit 12.3. + * @endrst + */ + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeysCopy( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_input_keys, + KeyIteratorT d_output_keys, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return SortKeysCopyNoNVTX( + d_temp_storage, temp_storage_bytes, d_input_keys, d_output_keys, num_items, compare_op, stream); + } + + //! @rst + //! Stably sorts keys using a merge sorting method. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - StableSortKeysCopy preserves the relative ordering of equivalent elements. + //! - Input array ``d_input_keys`` is not modified. + //! - The behavior is undefined if the input and output ranges overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_merge_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-keys-copy-env + //! :end-before: example-end stable-sort-keys-copy-env + //! + //! @endrst + //! + //! @tparam KeyInputIteratorT + //! **[inferred]** Random-access iterator type for input keys @iterator + //! + //! @tparam KeyIteratorT + //! **[inferred]** Random-access iterator type for output keys @iterator + //! + //! @tparam OffsetT + //! **[inferred]** Integer type for offsets + //! + //! @tparam CompareOpT + //! **[inferred]** Comparison function object type + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type + //! + //! @param[in] d_input_keys + //! Pointer to the input sequence of unsorted input keys + //! + //! @param[out] d_output_keys + //! Pointer to the output sequence of sorted input keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] compare_op + //! Comparison function object + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t StableSortKeysCopy( + KeyInputIteratorT d_input_keys, + KeyIteratorT d_output_keys, + OffsetT num_items, + CompareOpT compare_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using ChooseOffsetT = detail::choose_offset_t; + using default_policy_selector = detail::merge_sort::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::merge_sort::dispatch( + storage, + bytes, + d_input_keys, + static_cast(nullptr), + d_output_keys, + static_cast(nullptr), + static_cast(num_items), + compare_op, + stream, + policy_selector); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_partition.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_partition.cuh new file mode 100644 index 00000000..2c1aa8c4 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_partition.cuh @@ -0,0 +1,1060 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DevicePartition provides device-wide, parallel operations for partitioning sequences of data items residing +//! within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DevicePartition provides device-wide, parallel operations for +//! partitioning sequences of data items residing within device-accessible memory. +//! +//! Overview +//! ++++++++++++++++++++++++++ +//! +//! These operations apply a selection criterion to construct a partitioned +//! output sequence from items selected/unselected from a specified input +//! sequence. +//! +//! Usage Considerations +//! ++++++++++++++++++++++++++ +//! +//! @cdp_class{DevicePartition} +//! +//! Performance +//! ++++++++++++++++++++++++++ +//! +//! @linear_performance{partition} +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DevicePartition, except @p If with three partitions, that accept an environment can be tuned by +//! passing a custom :ref:`policy selector ` that returns a :cpp:struct:`cub::PartitionPolicy`, as +//! shown in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_partition_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin partition-if-policy-selector +//! :end-before: example-end partition-if-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_partition_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin partition-if-tuning +//! :end-before: example-end partition-if-tuning +//! +//! The environment overload of the three-way @p If algorithm can be tuned using a +//! :cpp:struct:`cub::ThreeWayPartitionPolicy` instead: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_partition_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin partition-three-way-policy-selector +//! :end-before: example-end partition-three-way-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_partition_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin partition-three-way-tuning +//! :end-before: example-end partition-three-way-tuning +//! +//! @endrst +struct DevicePartition +{ +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // Several algorithms dispatch to DeviceSelect, but we want to have a dedicated PartitionPolicy, so we need to adapt + // the policy selector to convert the tuning policy + template + struct __policy_selector_adapter + { + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> SelectPolicy + { + // the user-provided policy selector returns a PartitionPolicy, the default one a SelectPolicy + using policy_t = ::cuda::std::remove_cvref_t; + static_assert( + ::cuda::std::is_same_v || ::cuda::std::is_same_v); + const auto policy = PolicySelector{}(cc); + return SelectPolicy{ + SelectAlgorithm::lookback, + {policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + policy.lookback.lookback_delay}}; + } + }; +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Uses the ``d_flags`` sequence to split the corresponding items from + //! ``d_in`` into a partitioned sequence ``d_out``. + //! The total number of items copied into the first partition is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The value type of ``d_flags`` must be castable to ``bool`` (e.g., ``bool``, ``char``, ``int``, etc.). + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering, however copies of the unselected + //! items are compacted into the rear of ``d_out`` in reverse order. + //! - The range ``[d_out, d_out + num_items)`` shall not overlap + //! ``[d_in, d_in + num_items)`` nor ``[d_flags, d_flags + num_items)`` in any way. + //! The range ``[d_in, d_in + num_items)`` may overlap ``[d_flags, d_flags + num_items)``. + //! - @devicestorage + //! + //! Snippet + //! ++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input, flags, and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [1, 2, 3, 4, 5, 6, 7, 8] + //! char *d_flags; // e.g., [1, 0, 0, 1, 0, 1, 1, 0] + //! int *d_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DevicePartition::Flagged( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_flags, d_out, d_num_selected_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DevicePartition::Flagged( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_flags, d_out, d_num_selected_out, num_items); + //! + //! // d_out <-- [1, 4, 6, 7, 8, 5, 3, 2] + //! // d_num_selected_out <-- [4] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_out + //! Pointer to the output sequence of partitioned data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected (i.e., the + //! offset of the unselected partition) + //! + //! @param[in] num_items + //! Total number of items to select from + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Flagged( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FlagIterator d_flags, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DevicePartition::Flagged"); + using choose_offset_t = detail::choose_signed_offset; + using offset_t = typename choose_offset_t::type; + + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const auto error = choose_offset_t::is_exceeding_offset_type(num_items)) + { + return error; + } + + // we can't use dispatch_with_env_and_tuning, since default_policy_selector uses SelectPolicy, not PartitionPolicy + return detail::dispatch_with_env( + d_temp_storage, + temp_storage_bytes, + env, + [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + using default_policy_selector = detail::select:: + policy_selector_from_types; + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + return detail::select::dispatch( + storage, + bytes, + d_in, + d_flags, + d_out, + d_num_selected_out, + NullType{}, + NullType{}, + static_cast(num_items), + stream, + __policy_selector_adapter{}); + }); + } + + //! @rst + //! Uses the ``d_flags`` sequence to split the corresponding items from + //! ``d_in`` into a partitioned sequence ``d_out``. + //! The total number of items copied into the first partition is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The value type of ``d_flags`` must be castable to ``bool`` (e.g., ``bool``, ``char``, ``int``, etc.). + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering, however copies of the unselected + //! items are compacted into the rear of ``d_out`` in reverse order. + //! - The range ``[d_out, d_out + num_items)`` shall not overlap + //! ``[d_in, d_in + num_items)`` nor ``[d_flags, d_flags + num_items)`` in any way. + //! The range ``[d_in, d_in + num_items)`` may overlap ``[d_flags, d_flags + num_items)``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the partitioning of flagged items from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_partition_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin partition-flagged-env + //! :end-before: example-end partition-flagged-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_out + //! Pointer to the output sequence of partitioned data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected (i.e., the + //! offset of the unselected partition) + //! + //! @param[in] num_items + //! Total number of items to select from + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Flagged( + InputIteratorT d_in, + FlagIterator d_flags, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DevicePartition::Flagged"); + + using choose_offset_t = detail::choose_signed_offset; + using offset_t = typename choose_offset_t::type; + + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const auto error = choose_offset_t::is_exceeding_offset_type(num_items)) + { + return error; + } + + // we can't use dispatch_with_env_and_tuning, since default_policy_selector uses SelectPolicy, not PartitionPolicy + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + using default_policy_selector = detail::select:: + policy_selector_from_types; + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + return detail::select::dispatch( + storage, + bytes, + d_in, + d_flags, + d_out, + d_num_selected_out, + NullType{}, + NullType{}, + static_cast(num_items), + stream, + __policy_selector_adapter{}); + }); + } + + //! @rst + //! Uses the ``select_op`` functor to split the corresponding items from ``d_in`` into + //! a partitioned sequence ``d_out``. The total number of items copied into the first partition is written + //! to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering, however copies of the unselected + //! items are compacted into the rear of ``d_out`` in reverse order. + //! - The range ``[d_out, d_out + num_items)`` shall not overlap + //! ``[d_in, d_in + num_items)`` in any way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Functor type for selecting values less than some criteria + //! struct LessThan + //! { + //! int compare; + //! + //! CUB_RUNTIME_FUNCTION __forceinline__ + //! explicit LessThan(int compare) : compare(compare) {} + //! + //! CUB_RUNTIME_FUNCTION __forceinline__ + //! bool operator()(const int &a) const + //! { + //! return (a < compare); + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [0, 2, 3, 9, 5, 2, 81, 8] + //! int *d_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ ] + //! LessThan select_op(7); + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DevicePartition::If( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, d_num_selected_out, num_items, select_op); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DevicePartition::If( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, d_num_selected_out, num_items, select_op); + //! + //! // d_out <-- [0, 2, 3, 5, 2, 8, 81, 9] + //! // d_num_selected_out <-- [5] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection functor type having member `bool operator()(const T &a)` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of partitioned data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected (i.e., the offset of the unselected partition) + //! + //! @param[in] num_items + //! Total number of items to select from + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + If(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DevicePartition::If"); + + using choose_offset_t = detail::choose_signed_offset; + using offset_t = typename choose_offset_t::type; + + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const auto error = choose_offset_t::is_exceeding_offset_type(num_items)) + { + return error; + } + + return detail::dispatch_with_env( + d_temp_storage, + temp_storage_bytes, + env, + [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + using default_policy_selector = detail::select:: + policy_selector_from_types; + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + return detail::select::dispatch( + storage, + bytes, + d_in, + static_cast(nullptr), + d_out, + d_num_selected_out, + select_op, + NullType{}, + static_cast(num_items), + stream, + __policy_selector_adapter{}); + }); + } + + //! @rst + //! Uses the ``select_op`` functor to split the corresponding items from ``d_in`` into + //! a partitioned sequence ``d_out``. The total number of items copied into the first partition is written + //! to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering, however copies of the unselected + //! items are compacted into the rear of ``d_out`` in reverse order. + //! - The range ``[d_out, d_out + num_items)`` shall not overlap + //! ``[d_in, d_in + num_items)`` in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the partitioning of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_partition_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin partition-if-env + //! :end-before: example-end partition-if-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection functor type having member `bool operator()(const T &a)` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of partitioned data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected (i.e., the offset of the unselected partition) + //! + //! @param[in] num_items + //! Total number of items to select from + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + If(InputIteratorT d_in, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DevicePartition::If"); + + using choose_offset_t = detail::choose_signed_offset; + using offset_t = typename choose_offset_t::type; + + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const auto error = choose_offset_t::is_exceeding_offset_type(num_items)) + { + return error; + } + + // we can't use dispatch_with_env_and_tuning, since default_policy_selector uses SelectPolicy, not PartitionPolicy + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning_env, void* storage, size_t& bytes, cudaStream_t stream) { + using default_policy_selector = detail::select:: + policy_selector_from_types; + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + return detail::select::dispatch( + storage, + bytes, + d_in, + static_cast(nullptr), + d_out, + d_num_selected_out, + select_op, + NullType{}, + static_cast(num_items), + stream, + __policy_selector_adapter{}); + }); + } + + //! @rst + //! Uses two functors to split the corresponding items from ``d_in`` into a three partitioned sequences + //! ``d_first_part_out``, ``d_second_part_out``, and ``d_unselected_out``. + //! The total number of items copied into the first partition is written + //! to ``d_num_selected_out[0]``, while the total number of items copied into the second partition is written + //! to ``d_num_selected_out[1]``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Copies of the items selected by ``select_first_part_op`` are compacted + //! into ``d_first_part_out`` and maintain their original relative ordering. + //! - Copies of the items selected by ``select_second_part_op`` are compacted + //! into ``d_second_part_out`` and maintain their original relative ordering. + //! - Copies of the unselected items are compacted into the ``d_unselected_out`` in reverse order. + //! - The ranges ``[d_out, d_out + num_items)``, + //! ``[d_first_part_out, d_first_part_out + d_num_selected_out[0])``, + //! ``[d_second_part_out, d_second_part_out + d_num_selected_out[1])``, + //! ``[d_unselected_out, d_unselected_out + num_items - d_num_selected_out[0] - d_num_selected_out[1])``, + //! shall not overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates how this algorithm can partition an + //! input vector into small, medium, and large items so that the relative + //! order of items remain deterministic. + //! + //! Let's consider any value that doesn't exceed six a small one. On the + //! other hand, any value that exceeds 50 will be considered a large one. + //! Since the value used to define a small part doesn't match one that + //! defines the large part, the intermediate segment is implied. + //! + //! These definitions partition a value space into three categories. We want + //! to preserve the order of items in which they appear in the input vector. + //! Since the algorithm provides stable partitioning, this is possible. + //! + //! Since the number of items in each category is unknown beforehand, we need + //! three output arrays of num_items elements each. To reduce the memory + //! requirements, we can combine the output storage for two categories. + //! + //! Since each value falls precisely in one category, it's safe to add + //! "large" values into the head of the shared output vector and the "middle" + //! values into its tail. To add items into the tail of the output array, we + //! can use ``cuda::std::reverse_iterator``. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Functor type for selecting values less than some criteria + //! struct LessThan + //! { + //! int compare; + //! + //! __host__ __device__ __forceinline__ + //! explicit LessThan(int compare) : compare(compare) {} + //! + //! __host__ __device__ __forceinline__ + //! bool operator()(const int &a) const + //! { + //! return a < compare; + //! } + //! }; + //! + //! // Functor type for selecting values greater than some criteria + //! struct GreaterThan + //! { + //! int compare; + //! + //! __host__ __device__ __forceinline__ + //! explicit GreaterThan(int compare) : compare(compare) {} + //! + //! __host__ __device__ __forceinline__ + //! bool operator()(const int &a) const + //! { + //! return a > compare; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [0, 2, 3, 9, 5, 2, 81, 8] + //! int *d_large_and_unselected_out; // e.g., [ , , , , , , , ] + //! int *d_small_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ , ] + //! cud::std::reverse_iterator unselected_out(d_large_and_unselected_out + num_items); + //! LessThan small_items_selector(7); + //! GreaterThan large_items_selector(50); + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DevicePartition::If( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_large_and_medium_out, d_small_out, unselected_out, + //! d_num_selected_out, num_items, + //! large_items_selector, small_items_selector); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DevicePartition::If( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_large_and_medium_out, d_small_out, unselected_out, + //! d_num_selected_out, num_items, + //! large_items_selector, small_items_selector); + //! + //! // d_large_and_unselected_out <-- [ 81, , , , , , 8, 9 ] + //! // d_small_out <-- [ 0, 2, 3, 5, 2, , , ] + //! // d_num_selected_out <-- [ 1, 5 ] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam FirstOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output + //! items selected by first operator @iterator + //! + //! @tparam SecondOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output + //! items selected by second operator @iterator + //! + //! @tparam UnselectedOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing + //! unselected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items + //! selected @iterator + //! + //! @tparam SelectFirstPartOp + //! **[inferred]** Selection functor type having member `bool operator()(const T &a)` + //! + //! @tparam SelectSecondPartOp + //! **[inferred]** Selection functor type having member `bool operator()(const T &a)` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_first_part_out + //! Pointer to the output sequence of data items selected by `select_first_part_op` + //! + //! @param[out] d_second_part_out + //! Pointer to the output sequence of data items selected by `select_second_part_op` + //! + //! @param[out] d_unselected_out + //! Pointer to the output sequence of unselected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output array with two elements, where total number of + //! items selected by `select_first_part_op` is stored as + //! `d_num_selected_out[0]` and total number of items selected by + //! `select_second_part_op` is stored as `d_num_selected_out[1]`, + //! respectively + //! + //! @param[in] num_items + //! Total number of items to select from + //! + //! @param[in] select_first_part_op + //! Unary selection operator to select `d_first_part_out` + //! + //! @param[in] select_second_part_op + //! Unary selection operator to select `d_second_part_out` + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + If(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FirstOutputIteratorT d_first_part_out, + SecondOutputIteratorT d_second_part_out, + UnselectedOutputIteratorT d_unselected_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + SelectFirstPartOp select_first_part_op, + SelectSecondPartOp select_second_part_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DevicePartition::If"); + using choose_offset_t = detail::choose_signed_offset; + + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const auto error = choose_offset_t::is_exceeding_offset_type(num_items)) + { + return error; + } + + using offset_t = typename choose_offset_t::type; + using default_policy_selector = + detail::three_way_partition::policy_selector_from_types, + detail::three_way_partition::per_partition_offset_t>; + + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::three_way_partition::dispatch( + storage, + bytes, + d_in, + d_first_part_out, + d_second_part_out, + d_unselected_out, + d_num_selected_out, + select_first_part_op, + select_second_part_op, + static_cast(num_items), + stream, + policy_selector); + }); + } + + //! @rst + //! Uses two functors to split the corresponding items from ``d_in`` into three partitioned sequences + //! ``d_first_part_out``, ``d_second_part_out``, and ``d_unselected_out``. + //! The total number of items copied into the first partition is written + //! to ``d_num_selected_out[0]``, while the total number of items copied into the second partition is written + //! to ``d_num_selected_out[1]``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Copies of the items selected by ``select_first_part_op`` are compacted + //! into ``d_first_part_out`` and maintain their original relative ordering. + //! - Copies of the items selected by ``select_second_part_op`` are compacted + //! into ``d_second_part_out`` and maintain their original relative ordering. + //! - Copies of the unselected items are compacted into the ``d_unselected_out`` in reverse order. + //! - The ranges ``[d_out, d_out + num_items)``, + //! ``[d_first_part_out, d_first_part_out + d_num_selected_out[0])``, + //! ``[d_second_part_out, d_second_part_out + d_num_selected_out[1])``, + //! ``[d_unselected_out, d_unselected_out + num_items - d_num_selected_out[0] - d_num_selected_out[1])``, + //! shall not overlap in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates three-way partitioning. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_partition_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin partition-three-way-env + //! :end-before: example-end partition-three-way-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam FirstOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output + //! items selected by first operator @iterator + //! + //! @tparam SecondOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output + //! items selected by second operator @iterator + //! + //! @tparam UnselectedOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing + //! unselected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items + //! selected @iterator + //! + //! @tparam SelectFirstPartOp + //! **[inferred]** Selection functor type having member `bool operator()(const T &a)` + //! + //! @tparam SelectSecondPartOp + //! **[inferred]** Selection functor type having member `bool operator()(const T &a)` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_first_part_out + //! Pointer to the output sequence of data items selected by `select_first_part_op` + //! + //! @param[out] d_second_part_out + //! Pointer to the output sequence of data items selected by `select_second_part_op` + //! + //! @param[out] d_unselected_out + //! Pointer to the output sequence of unselected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output array with two elements, where total number of + //! items selected by `select_first_part_op` is stored as + //! `d_num_selected_out[0]` and total number of items selected by + //! `select_second_part_op` is stored as `d_num_selected_out[1]`, + //! respectively + //! + //! @param[in] num_items + //! Total number of items to select from + //! + //! @param[in] select_first_part_op + //! Unary selection operator to select `d_first_part_out` + //! + //! @param[in] select_second_part_op + //! Unary selection operator to select `d_second_part_out` + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + If(InputIteratorT d_in, + FirstOutputIteratorT d_first_part_out, + SecondOutputIteratorT d_second_part_out, + UnselectedOutputIteratorT d_unselected_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + SelectFirstPartOp select_first_part_op, + SelectSecondPartOp select_second_part_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DevicePartition::If"); + + using choose_offset_t = detail::choose_signed_offset; + if (const auto error = choose_offset_t::is_exceeding_offset_type(num_items)) + { + return error; + } + + using offset_t = typename choose_offset_t::type; + using default_policy_selector = + detail::three_way_partition::policy_selector_from_types, + detail::three_way_partition::per_partition_offset_t>; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::three_way_partition::dispatch( + storage, + bytes, + d_in, + d_first_part_out, + d_second_part_out, + d_unselected_out, + d_num_selected_out, + select_first_part_op, + select_second_part_op, + static_cast(num_items), + stream, + policy_selector); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_radix_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_radix_sort.cuh new file mode 100644 index 00000000..cb105ee5 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_radix_sort.cuh @@ -0,0 +1,5272 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceRadixSort provides device-wide, parallel operations for computing a radix sort across a sequence of data +//! items residing within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document +template > +inline constexpr bool __can_use_radix_sort = + (::cuda::std::is_arithmetic_v<_ValueType> +# if _CCCL_HAS_NVFP16() && !defined(__CUDA_NO_HALF_OPERATORS__) && !defined(__CUDA_NO_HALF_CONVERSIONS__) + || ::cuda::std::is_same_v<_ValueType, __half> +# endif // _CCCL_HAS_NVFP16() && !defined(__CUDA_NO_HALF_OPERATORS__) && !defined(__CUDA_NO_HALF_CONVERSIONS__) +# if _CCCL_HAS_NVBF16() && !defined(__CUDA_NO_BFLOAT16_CONVERSIONS__) && !defined(__CUDA_NO_BFLOAT16_OPERATORS__) + || ::cuda::std::is_same_v<_ValueType, __nv_bfloat16> +# endif // _CCCL_HAS_NVBF16() && !defined(__CUDA_NO_BFLOAT16_CONVERSIONS__) && + // !defined(__CUDA_NO_BFLOAT16_OPERATORS__) + ) + && ::cuda::std::__is_one_of_v<::cuda::std::remove_cvref_t<_BinaryPredicate>, + ::cuda::std::less<>, + ::cuda::std::less<_ValueType>, + ::cuda::std::greater<>, + ::cuda::std::greater<_ValueType>>; +#endif // !_CCCL_DOXYGEN_INVOKED + +//! @rst +//! DeviceRadixSort provides device-wide, parallel operations for +//! computing a radix sort across a sequence of data items residing +//! within device-accessible memory. +//! +//! .. image:: ../../img/sorting_logo.png +//! :align: center +//! +//! Overview +//! -------------------------------------------------- +//! +//! The `radix sorting method `_ +//! arranges items into ascending (or descending) order. The algorithm relies +//! upon a positional representation for keys, i.e., each key is comprised of an +//! ordered sequence of symbols (e.g., digits, characters, etc.) specified from +//! least-significant to most-significant. For a given input sequence of keys +//! and a set of rules specifying a total ordering of the symbolic alphabet, the +//! radix sorting method produces a lexicographic ordering of those keys. +//! +//! @rowmajor +//! +//! Supported Types +//! -------------------------------------------------- +//! +//! DeviceRadixSort can sort 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. +//! +//! Floating-Point Special Cases +//! -------------------------------------------------- +//! +//! - Positive and negative zeros are considered equivalent, and will be treated +//! as such in the output. +//! - No special handling is implemented for NaN values; these are sorted +//! according to their bit representations after any transformations. +//! +//! Transformations +//! -------------------------------------------------- +//! +//! Although the direct radix sorting method can only be applied to unsigned +//! integral types, DeviceRadixSort is able to sort signed and floating-point +//! types via simple bit-wise transformations that ensure lexicographic key +//! ordering. Additional transformations occur for descending sorts. These +//! transformations must be considered when restricting the +//! ``[begin_bit, end_bit)`` range, as the bitwise transformations will occur +//! before the bit-range truncation. +//! +//! Any transformations applied to the keys prior to sorting are reversed +//! while writing to the final output buffer. +//! +//! Type Specific Bitwise Transformations +//! -------------------------------------------------- +//! +//! To convert the input values into a radix-sortable bitwise representation, +//! the following transformations take place prior to sorting: +//! +//! - For unsigned integral values, the keys are used directly. +//! - For signed integral values, the sign bit is inverted. +//! - For positive floating point values, the sign bit is inverted. +//! - For negative floating point values, the full key is inverted. +//! +//! For floating point types, positive and negative zero are a special case and +//! will be considered equivalent during sorting. +//! +//! Descending Sort Bitwise Transformations +//! -------------------------------------------------- +//! +//! If descending sort is used, the keys are inverted after performing any +//! type-specific transformations, and the resulting keys are sorted in ascending +//! order. +//! +//! Stability +//! -------------------------------------------------- +//! +//! DeviceRadixSort is stable. For floating-point types, ``-0.0`` and ``+0.0`` are +//! considered equal and appear in the result in the same order as they appear in +//! the input. +//! +//! Usage Considerations +//! -------------------------------------------------- +//! +//! @cdp_class{DeviceRadixSort} +//! +//! Performance +//! -------------------------------------------------- +//! +//! @linear_performance{radix sort} +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceRadixSort that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::RadixSortPolicy`, as shown in the +//! example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin radix-sort-keys-policy-selector +//! :end-before: example-end radix-sort-keys-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin radix-sort-keys-tuning +//! :end-before: example-end radix-sort-keys-tuning +//! +//! @endrst +struct DeviceRadixSort +{ +private: + // TODO(bgruber): I would ideally like to have the logic of extracting the policy selector from the tuning environment + // inside the dispatch function, but this will not work with CCCL.C, which needs to pass a stateful policy selector. + // Refactor this once we have a host code JIT compiler. + template > + CUB_RUNTIME_FUNCTION static cudaError_t select_tuning_and_dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + OffsetT num_items, + int begin_bit, + int end_bit, + bool is_overwrite_okay, + cudaStream_t stream, + DecomposerT decomposer = {}, + TuningEnvT = {}) + { + using default_policy_selector_t = detail::radix_sort::policy_selector_from_types; + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + return detail::radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + begin_bit, + end_bit, + is_overwrite_okay, + stream, + decomposer, + policy_selector_t{}); + } + + template > + CUB_RUNTIME_FUNCTION static cudaError_t radix_sort_with_decomposer( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream, + TuningEnvT tuning_env = {}, + int begin_bit = 0, + int end_bit = detail::radix::traits_t::default_end_bit(DecomposerT{}), + bool is_overwrite_okay = true) + { + using offset_t = detail::choose_offset_t; + static constexpr bool decomposer_check = detail::radix::decomposer_check; + + static_assert(decomposer_check, + "DecomposerT must be a callable object returning a tuple of references to " + "arithmetic types"); + + if constexpr (decomposer_check) + { + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream, + decomposer, + tuning_env); + } + _CCCL_UNREACHABLE(); + } + + template > + CUB_RUNTIME_FUNCTION static cudaError_t radix_sort_with_decomposer( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream, + TuningEnvT tuning_env = {}, + int begin_bit = 0, + int end_bit = detail::radix::traits_t::default_end_bit(DecomposerT{})) + { + // We cast away const-ness, but will *not* write to these arrays. ``DispatchRadixSort::Dispatch`` will allocate + // temporary storage and create a new double-buffer internally when the ``is_overwrite_ok`` flag is not set. + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + decomposer, + stream, + tuning_env, + begin_bit, + end_bit, + /* is_overwrite_okay */ false); + } + + // Name reported for NVTX ranges + _CCCL_HOST_DEVICE static constexpr auto GetName() -> const char* + { + return "cub::DeviceRadixSort"; + } + +public: + //! @name KeyT-value pairs + //! @{ + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + num_items)`` + //! - ``[d_values_in, d_values_in + num_items)`` + //! - ``[d_values_out, d_values_out + num_items)`` + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageNP For sorting using only ``O(P)`` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the sorting of a device vector of ``int`` + //! keys with associated vector of ``int`` values. + //! @endrst + //! + //! @code{.cpp} + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [ ... ] + //! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_values_out; // e.g., [ ... ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, num_items); + //! + //! // d_keys_out <-- [0, 3, 5, 6, 7, 8, 9] + //! // d_values_out <-- [5, 4, 3, 1, 2, 0, 6] + //! @endcode + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + // Unsigned integer type for global offsets. + using OffsetT = detail::choose_offset_t; + + // TODO API that doesn't accept decomposer should also contain a static + // assert that the key type is fundamental. + + // We cast away const-ness, but will *not* write to these arrays. + // ``DispatchRadixSort::Dispatch`` will allocate temporary storage and + // create a new double-buffer internally when the ``is_overwrite_ok`` flag + // is not set. + constexpr bool is_overwrite_okay = false; + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + num_items)`` + //! - ``[d_values_in, d_values_in + num_items)`` + //! - ``[d_values_out, d_values_out + num_items)`` + //! + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the env-based sorting of key-value pairs: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-env + //! :end-before: example-end radix-sort-pairs-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_keys_in + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using offset_t = detail::choose_offset_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + tuning_env); + }); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! * The contents of the input data are not altered by the sorting operation. + //! * Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys_in, d_keys_in + num_items)`` + //! * ``[d_keys_out, d_keys_out + num_items)`` + //! * ``[d_values_in, d_values_in + num_items)`` + //! * ``[d_values_out, d_values_out + num_items)`` + //! + //! * A bit subrange ``[begin_bit, end_bit)`` is provided to specify + //! differentiating key bits. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! * @devicestorageNP For sorting using only :math:`O(P)` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortPairs``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-bits + //! :end-before: example-end pairs-bits + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``(sizeof(float) + sizeof(long long int)) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + decomposer, + stream, + {}, + begin_bit, + end_bit); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - A bit subrange ``[begin_bit, end_bit)`` is provided to specify which + //! key bits are used for sorting. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-decomposer-bits-env + //! :end-before: example-end radix-sort-pairs-decomposer-bits-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types + //! + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in] d_keys_in + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, + bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + decomposer, + stream, + tuning_env, + begin_bit, + end_bit); + }); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The contents of the input data are not altered by the sorting operation. + //! * Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys_in, d_keys_in + num_items)`` + //! * ``[d_keys_out, d_keys_out + num_items)`` + //! * ``[d_values_in, d_values_in + num_items)`` + //! * ``[d_values_out, d_values_out + num_items)`` + //! + //! * @devicestorageNP For sorting using only :math:`O(P)` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortPairs``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs + //! :end-before: example-end pairs + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + decomposer, + stream); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-decomposer-env + //! :end-before: example-end radix-sort-pairs-decomposer-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types + //! + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in] d_keys_in + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, num_items, decomposer, stream, tuning_env); + }); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! - ``[d_values.Current(), d_values.Current() + num_items)`` + //! - ``[d_values.Alternate(), d_values.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the sorting of a device vector of ``int`` + //! keys with associated vector of ``int`` values. + //! @endrst + //! + //! @code + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // sorting data + //! int num_items; // e.g., 7 + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [ ... ] + //! int *d_value_buf; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_value_alt_buf; // e.g., [ ... ] + //! ... + //! + //! // Create a set of DoubleBuffers to wrap pairs of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! cub::DoubleBuffer d_values(d_value_buf, d_value_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRadixSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceRadixSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items); + //! + //! // d_keys.Current() <-- [0, 3, 5, 6, 7, 8, 9] + //! // d_values.Current() <-- [5, 4, 3, 1, 2, 0, 6] + //! + //! @endcode + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Unsigned integer type for global offsets. + using OffsetT = detail::choose_offset_t; + + constexpr bool is_overwrite_okay = true; + + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! - ``[d_values.Current(), d_values.Current() + num_items)`` + //! - ``[d_values.Alternate(), d_values.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using offset_t = detail::choose_offset_t; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + tuning_env); + }); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! * The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! - ``[d_values.Current(), d_values.Current() + num_items)`` + //! - ``[d_values.Alternate(), d_values.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortPairs``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-db + //! :end-before: example-end pairs-db + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-db-decomposer-env + //! :end-before: example-end radix-sort-pairs-db-decomposer-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types + //! + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys + //! + //! @param[in,out] d_values + //! Double-buffer of values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env); + }); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! * The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! - ``[d_values.Current(), d_values.Current() + num_items)`` + //! - ``[d_values.Alternate(), d_values.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortPairs``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-bits-db + //! :end-before: example-end pairs-bits-db + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``(sizeof(float) + sizeof(long long int)) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, {}, begin_bit, end_bit); + } + + //! @rst + //! Sorts key-value pairs into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure. + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! - A bit subrange ``[begin_bit, end_bit)`` is provided to specify which + //! key bits are used for sorting. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-db-decomposer-bits-env + //! :end-before: example-end radix-sort-pairs-db-decomposer-bits-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types + //! + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys + //! + //! @param[in,out] d_values + //! Double-buffer of values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairs( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env, begin_bit, end_bit); + }); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + num_items)`` + //! - ``[d_values_in, d_values_in + num_items)`` + //! - ``[d_values_out, d_values_out + num_items)`` + //! + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageNP For sorting using only ``O(P)`` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the sorting of a device vector of ``int`` + //! keys with associated vector of ``int`` values. + //! @endrst + //! + //! @code{.cpp} + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [ ... ] + //! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_values_out; // e.g., [ ... ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRadixSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceRadixSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, num_items); + //! + //! // d_keys_out <-- [9, 8, 7, 6, 5, 3, 0] + //! // d_values_out <-- [6, 0, 2, 1, 3, 4, 5] + //! @endcode + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Unsigned integer type for global offsets. + using OffsetT = detail::choose_offset_t; + + // We cast away const-ness, but will *not* write to these arrays. + // ``DispatchRadixSort::Dispatch`` will allocate temporary storage and + // create a new double-buffer internally when the ``is_overwrite_ok`` flag + // is not set. + constexpr bool is_overwrite_okay = false; + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + num_items)`` + //! - ``[d_values_in, d_values_in + num_items)`` + //! - ``[d_values_out, d_values_out + num_items)`` + //! + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the env-based descending sort of key-value pairs: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-descending-env + //! :end-before: example-end radix-sort-pairs-descending-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_keys_in + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using offset_t = detail::choose_offset_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + tuning_env); + }); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! * The contents of the input data are not altered by the sorting operation. + //! * Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys_in, d_keys_in + num_items)`` + //! * ``[d_keys_out, d_keys_out + num_items)`` + //! * ``[d_values_in, d_values_in + num_items)`` + //! * ``[d_values_out, d_values_out + num_items)`` + //! + //! * A bit subrange ``[begin_bit, end_bit)`` is provided to specify + //! differentiating key bits. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! * @devicestorageNP For sorting using only :math:`O(P)` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortPairsDescending``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-descending-bits + //! :end-before: example-end pairs-descending-bits + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``(sizeof(float) + sizeof(long long int)) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + decomposer, + stream, + {}, + begin_bit, + end_bit); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The contents of the input data are not altered by the sorting operation. + //! * Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys_in, d_keys_in + num_items)`` + //! * ``[d_keys_out, d_keys_out + num_items)`` + //! * ``[d_values_in, d_values_in + num_items)`` + //! * ``[d_values_out, d_values_out + num_items)`` + //! + //! * @devicestorageNP For sorting using only :math:`O(P)` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortPairsDescending``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-descending + //! :end-before: example-end pairs-descending + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Pointer to the correspondingly-reordered output sequence of associated + //! value items + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + decomposer, + stream); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! - ``[d_values.Current(), d_values.Current() + num_items)`` + //! - ``[d_values.Alternate(), d_values.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the sorting of a device vector of ``int`` + //! keys with associated vector of ``int`` values. + //! @endrst + //! + //! @code{.cpp} + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [ ... ] + //! int *d_value_buf; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_value_alt_buf; // e.g., [ ... ] + //! ... + //! + //! // Create a set of DoubleBuffers to wrap pairs of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! cub::DoubleBuffer d_values(d_value_buf, d_value_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRadixSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceRadixSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items); + //! + //! // d_keys.Current() <-- [9, 8, 7, 6, 5, 3, 0] + //! // d_values.Current() <-- [6, 0, 2, 1, 3, 4, 5] + //! @endcode + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Unsigned integer type for global offsets. + using OffsetT = detail::choose_offset_t; + + constexpr bool is_overwrite_okay = true; + + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! - ``[d_values.Current(), d_values.Current() + num_items)`` + //! - ``[d_values.Alternate(), d_values.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using offset_t = detail::choose_offset_t; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + tuning_env); + }); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! * The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! - ``[d_values.Current(), d_values.Current() + num_items)`` + //! - ``[d_values.Alternate(), d_values.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortPairsDescending``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-descending-db + //! :end-before: example-end pairs-descending-db + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! * The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! - ``[d_values.Current(), d_values.Current() + num_items)`` + //! - ``[d_values.Alternate(), d_values.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortPairsDescending``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin pairs-descending-bits-db + //! :end-before: example-end pairs-descending-bits-db + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam ValueT + //! **[inferred]** ValueT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``(sizeof(float) + sizeof(long long int)) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, {}, begin_bit, end_bit); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - A bit subrange ``[begin_bit, end_bit)`` is provided to specify which key bits are used for sorting. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-descending-decomposer-bits-env + //! :end-before: example-end radix-sort-pairs-descending-decomposer-bits-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam ValueT + //! **[inferred]** ValueT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in] d_keys_in Pointer to input key data + //! @param[out] d_keys_out Pointer to sorted output key data + //! @param[in] d_values_in Pointer to input value data + //! @param[out] d_values_out Pointer to sorted output value data + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] begin_bit Least-significant bit index (inclusive) + //! @param[in] end_bit Most-significant bit index (exclusive) + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, + bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + decomposer, + stream, + tuning_env, + begin_bit, + end_bit); + }); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-descending-decomposer-env + //! :end-before: example-end radix-sort-pairs-descending-decomposer-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam ValueT + //! **[inferred]** ValueT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in] d_keys_in Pointer to input key data + //! @param[out] d_keys_out Pointer to sorted output key data + //! @param[in] d_values_in Pointer to input value data + //! @param[out] d_values_out Pointer to sorted output value data + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + NumItemsT num_items, + DecomposerT decomposer, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, bytes, d_keys_in, d_keys_out, d_values_in, d_values_out, num_items, decomposer, stream, tuning_env); + }); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given key and value buffers managed by DoubleBuffer structures. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-descending-db-decomposer-env + //! :end-before: example-end radix-sort-pairs-descending-db-decomposer-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam ValueT + //! **[inferred]** ValueT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in,out] d_keys Reference to the double-buffer of keys + //! @param[in,out] d_values Reference to the double-buffer of values + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env); + }); + } + + //! @rst + //! Sorts key-value pairs into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given key and value buffers managed by DoubleBuffer structures. + //! - A bit subrange ``[begin_bit, end_bit)`` is provided to specify which key bits are used for sorting. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-pairs-descending-db-decomposer-bits-env + //! :end-before: example-end radix-sort-pairs-descending-db-decomposer-bits-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam ValueT + //! **[inferred]** ValueT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in,out] d_keys Reference to the double-buffer of keys + //! @param[in,out] d_values Reference to the double-buffer of values + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] begin_bit Least-significant bit index (inclusive) + //! @param[in] end_bit Most-significant bit index (exclusive) + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortPairsDescending( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env, begin_bit, end_bit); + }); + } + + //! @} + //! @name Keys-only + //! @{ + + //! @rst + //! Sorts keys into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + num_items)`` + //! + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageNP For sorting using only ``O(P)`` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the sorting of a device vector of + //! ``int`` keys. + //! @endrst + //! + //! @code{.cpp} + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [ ... ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRadixSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceRadixSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, num_items); + //! + //! // d_keys_out <-- [0, 3, 5, 6, 7, 8, 9] + //! @endcode + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Unsigned integer type for global offsets. + using OffsetT = detail::choose_offset_t; + + // We cast away const-ness, but will *not* write to these arrays. + // ``DispatchRadixSort::Dispatch`` will allocate temporary storage and + // create a new double-buffer internally when the ``is_overwrite_ok`` flag + // is not set. + constexpr bool is_overwrite_okay = false; + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + // Null value type + DoubleBuffer d_values; + + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + num_items)`` + //! + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the env-based sorting of keys: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-env + //! :end-before: example-end radix-sort-keys-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_keys_in + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys( + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using offset_t = detail::choose_offset_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + tuning_env); + }); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The contents of the input data are not altered by the sorting operation. + //! * Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys_in, d_keys_in + num_items)`` + //! * ``[d_keys_out, d_keys_out + num_items)`` + //! + //! * A bit subrange ``[begin_bit, end_bit)`` is provided to specify + //! differentiating key bits. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! * @devicestorageNP For sorting using only :math:`O(P)` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortKeys``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-bits + //! :end-before: example-end keys-bits + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``(sizeof(float) + sizeof(long long int)) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + decomposer, + stream, + {}, + begin_bit, + end_bit); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - A bit subrange ``[begin_bit, end_bit)`` is provided to specify which key bits are used for sorting. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-custom-decomposer + //! :end-before: example-end radix-sort-keys-custom-decomposer + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-decomposer-bits-env + //! :end-before: example-end radix-sort-keys-decomposer-bits-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in] d_keys_in Pointer to input key data + //! @param[out] d_keys_out Pointer to sorted output key data + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] begin_bit Least-significant bit index (inclusive) + //! @param[in] end_bit Most-significant bit index (exclusive) + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys( + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + decomposer, + stream, + tuning_env, + begin_bit, + end_bit); + }); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The contents of the input data are not altered by the sorting operation. + //! * Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys_in, d_keys_in + num_items)`` + //! * ``[d_keys_out, d_keys_out + num_items)`` + //! + //! * An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! * @devicestorageNP For sorting using only :math:`O(P)` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortKeys``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys + //! :end-before: example-end keys + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + decomposer, + stream); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-custom-decomposer + //! :end-before: example-end radix-sort-keys-custom-decomposer + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-decomposer-env + //! :end-before: example-end radix-sort-keys-decomposer-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in] d_keys_in Pointer to input key data + //! @param[out] d_keys_out Pointer to sorted output key data + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + SortKeys(const KeyT* d_keys_in, KeyT* d_keys_out, NumItemsT num_items, DecomposerT decomposer, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + decomposer, + stream, + tuning_env); + }); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the sorting of a device vector of + //! ``int`` keys. + //! @endrst + //! + //! @code{.cpp} + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [ ... ] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRadixSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceRadixSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys, num_items); + //! + //! // d_keys.Current() <-- [0, 3, 5, 6, 7, 8, 9] + //! @endcode + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Unsigned integer type for global offsets. + using OffsetT = detail::choose_offset_t; + + constexpr bool is_overwrite_okay = true; + + // Null value type + DoubleBuffer d_values; + + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the env-based sorting of keys using DoubleBuffer: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-db-env + //! :end-before: example-end radix-sort-keys-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer stores the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys( + DoubleBuffer& d_keys, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using offset_t = detail::choose_offset_t; + + DoubleBuffer d_values; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + tuning_env); + }); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! * The contents of both buffers may be altered by the sorting operation. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! * ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! + //! * Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! * @devicestorageP + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortKeys``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-db + //! :end-before: example-end keys-db + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + DoubleBuffer d_values; + return radix_sort_with_decomposer( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a key buffer managed by a DoubleBuffer structure. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-custom-decomposer + //! :end-before: example-end radix-sort-keys-custom-decomposer + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-db-decomposer-env + //! :end-before: example-end radix-sort-keys-db-decomposer-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in,out] d_keys Reference to the double-buffer of keys + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + SortKeys(DoubleBuffer& d_keys, NumItemsT num_items, DecomposerT decomposer, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + DoubleBuffer d_values; + return radix_sort_with_decomposer( + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env); + }); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! * The contents of both buffers may be altered by the sorting operation. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! * ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! + //! * A bit subrange ``[begin_bit, end_bit)`` is provided to specify + //! differentiating key bits. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! * Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! * @devicestorageP + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortKeys``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-bits-db + //! :end-before: example-end keys-bits-db + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``(sizeof(float) + sizeof(long long int)) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + DoubleBuffer d_values; + return radix_sort_with_decomposer( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, {}, begin_bit, end_bit); + } + + //! @rst + //! Sorts keys into ascending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a key buffer managed by a DoubleBuffer structure. + //! - A bit subrange ``[begin_bit, end_bit)`` is provided to specify which key bits are used for sorting. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-custom-decomposer + //! :end-before: example-end radix-sort-keys-custom-decomposer + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-db-decomposer-bits-env + //! :end-before: example-end radix-sort-keys-db-decomposer-bits-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in,out] d_keys Reference to the double-buffer of keys + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] begin_bit Least-significant bit index (inclusive) + //! @param[in] end_bit Most-significant bit index (exclusive) + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeys( + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + DoubleBuffer d_values; + return radix_sort_with_decomposer( + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env, begin_bit, end_bit); + }); + } + + //! @rst Sorts keys into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + num_items)`` + //! + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageNP For sorting using only ``O(P)`` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the sorting of a device vector of + //! ``int`` keys. + //! @endrst + //! + //! @code{.cpp} + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [ ... ] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRadixSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceRadixSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, num_items); + //! + //! // d_keys_out <-- [9, 8, 7, 6, 5, 3, 0]s + //! @endcode + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Unsigned integer type for global offsets. + using OffsetT = detail::choose_offset_t; + + // We cast away const-ness, but will *not* write to these arrays. + // ``DispatchRadixSort::Dispatch`` will allocate temporary storage and + // create a new double-buffer internally when the ``is_overwrite_ok`` flag + // is not set. + constexpr bool is_overwrite_okay = false; + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values; + + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + num_items)`` + //! + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the env-based descending sort of keys: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-descending-env + //! :end-before: example-end radix-sort-keys-descending-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_keys_in + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using offset_t = detail::choose_offset_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + false, + stream, + {}, + tuning_env); + }); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The contents of the input data are not altered by the sorting operation. + //! * Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys_in, d_keys_in + num_items)`` + //! * ``[d_keys_out, d_keys_out + num_items)`` + //! + //! * An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! * @devicestorageNP For sorting using only :math:`O(P)` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortKeysDescending``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-descending-bits + //! :end-before: example-end keys-descending-bits + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``(sizeof(float) + sizeof(long long int)) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + decomposer, + stream, + {}, + begin_bit, + end_bit); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The contents of the input data are not altered by the sorting operation. + //! * Pointers to contiguous memory must be used; iterators are not currently + //! supported. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys_in, d_keys_in + num_items)`` + //! * ``[d_keys_out, d_keys_out + num_items)`` + //! + //! * @devicestorageNP For sorting using only :math:`O(P)` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortKeysDescending``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-descending + //! :end-before: example-end keys-descending + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @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 + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return radix_sort_with_decomposer( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + decomposer, + stream); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the sorting of a device vector of ``int`` keys. + //! @endrst + //! + //! @code{.cpp} + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [ ... ] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRadixSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceRadixSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, num_items); + //! + //! // d_keys.Current() <-- [9, 8, 7, 6, 5, 3, 0] + //! @endcode + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Unsigned integer type for global offsets. + using OffsetT = detail::choose_offset_t; + + constexpr bool is_overwrite_okay = true; + + // Null value type + DoubleBuffer d_values; + + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + is_overwrite_okay, + stream); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! - ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! The code snippet below illustrates the env-based descending sort of keys using DoubleBuffer: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-descending-db-env + //! :end-before: example-end radix-sort-keys-descending-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer stores the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param[in] begin_bit + //! The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( + DoubleBuffer& d_keys, + NumItemsT num_items, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using offset_t = detail::choose_offset_t; + + DoubleBuffer d_values; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return select_tuning_and_dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast(num_items), + begin_bit, + end_bit, + true, + stream, + {}, + tuning_env); + }); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! * The contents of both buffers may be altered by the sorting operation. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! * ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! + //! * Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! * @devicestorageP + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortKeysDescending``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-descending-db + //! :end-before: example-end keys-descending-db + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + DoubleBuffer d_values; + return radix_sort_with_decomposer( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! * The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! * The contents of both buffers may be altered by the sorting operation. + //! * In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! * ``[d_keys.Current(), d_keys.Current() + num_items)`` + //! * ``[d_keys.Alternate(), d_keys.Alternate() + num_items)`` + //! + //! * A bit subrange ``[begin_bit, end_bit)`` is provided to specify + //! differentiating key bits. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! * Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! * @devicestorageP + //! * @devicestorage + //! + //! Snippet + //! -------------------------------------------------- + //! + //! Let's consider a user-defined ``custom_t`` type below. To sort 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_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin custom-type + //! :end-before: example-end custom-type + //! + //! The following snippet shows how to sort an array of ``custom_t`` objects + //! using ``cub::DeviceRadixSort::SortKeysDescending``: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_custom.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin keys-descending-bits-db + //! :end-before: example-end keys-descending-bits-db + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam DecomposerT + //! **[inferred]** Type of a callable object responsible for decomposing a + //! ``KeyT`` into a tuple of references to its constituent arithmetic types: + //! ``cuda::std::tuple operator()(KeyT &key)``. + //! The leftmost element of the tuple is considered the most significant. + //! The call operator must not modify members of the key. + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of ``d_temp_storage`` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! Number of items to sort + //! + //! @param decomposer + //! Callable object responsible for decomposing a ``KeyT`` into a tuple of + //! references to its constituent arithmetic types. The leftmost element of + //! the tuple is considered the most significant. The call operator must not + //! modify members of the key. + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for + //! key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., ``(sizeof(float) + sizeof(long long int)) * 8``) + //! + //! @param[in] stream + //! **[optional]** CUDA stream to launch kernels within. + //! Default is stream0. + template , int> = 0> + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + DoubleBuffer d_values; + return radix_sort_with_decomposer( + d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items, decomposer, stream, {}, begin_bit, end_bit); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - A bit subrange ``[begin_bit, end_bit)`` is provided to specify which key bits are used for sorting. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-descending-decomposer-bits-env + //! :end-before: example-end radix-sort-keys-descending-decomposer-bits-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in] d_keys_in Pointer to input key data + //! @param[out] d_keys_out Pointer to sorted output key data + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] begin_bit Least-significant bit index (inclusive) + //! @param[in] end_bit Most-significant bit index (exclusive) + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + decomposer, + stream, + tuning_env, + begin_bit, + end_bit); + }); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx 2N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-descending-decomposer-env + //! :end-before: example-end radix-sort-keys-descending-decomposer-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in] d_keys_in Pointer to input key data + //! @param[out] d_keys_out Pointer to sorted output key data + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( + const KeyT* d_keys_in, KeyT* d_keys_out, NumItemsT num_items, DecomposerT decomposer, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return radix_sort_with_decomposer( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + decomposer, + stream, + tuning_env); + }); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a key buffer managed by a DoubleBuffer structure. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-descending-db-decomposer-env + //! :end-before: example-end radix-sort-keys-descending-db-decomposer-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in,out] d_keys Reference to the double-buffer of keys + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + SortKeysDescending(DoubleBuffer& d_keys, NumItemsT num_items, DecomposerT decomposer, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + DoubleBuffer d_values; + return radix_sort_with_decomposer( + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env); + }); + } + + //! @rst + //! Sorts keys into descending order using :math:`\approx N` auxiliary storage. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a key buffer managed by a DoubleBuffer structure. + //! - A bit subrange ``[begin_bit, end_bit)`` is provided to specify which key bits are used for sorting. + //! + //! Snippet + //! -------------------------------------------------- + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin radix-sort-keys-descending-db-decomposer-bits-env + //! :end-before: example-end radix-sort-keys-descending-db-decomposer-bits-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** KeyT type + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! @tparam DecomposerT + //! **[inferred]** Decomposer type + //! @tparam EnvT + //! **[inferred]** Environment type + //! + //! @param[in,out] d_keys Reference to the double-buffer of keys + //! @param[in] num_items Number of items to sort + //! @param[in] decomposer Decomposer callable + //! @param[in] begin_bit Least-significant bit index (inclusive) + //! @param[in] end_bit Most-significant bit index (exclusive) + //! @param[in] env **[optional]** Execution environment + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t SortKeysDescending( + DoubleBuffer& d_keys, + NumItemsT num_items, + DecomposerT decomposer, + int begin_bit, + int end_bit, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + DoubleBuffer d_values; + return radix_sort_with_decomposer( + storage, bytes, d_keys, d_values, num_items, decomposer, stream, tuning_env, begin_bit, end_bit); + }); + } + + //! @} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_reduce.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_reduce.cuh new file mode 100644 index 00000000..a6a94a05 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_reduce.cuh @@ -0,0 +1,2563 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceReduce provides device-wide, parallel operations for computing a reduction across a sequence of data +//! items residing within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +template +inline constexpr bool is_non_deterministic_v = + ::cuda::std::is_same_v; +} // namespace detail + +//! @rst +//! DeviceReduce provides device-wide, parallel operations for computing +//! a reduction across a sequence of data items residing within +//! device-accessible memory. +//! +//! .. image:: ../../img/reduce_logo.png +//! :align: center +//! +//! Overview +//! ==================================== +//! +//! A `reduction `_ +//! (or *fold*) uses a binary combining operator to compute a single aggregate +//! from a sequence of input elements. +//! +//! Usage Considerations +//! ==================================== +//! +//! @cdp_class{DeviceReduce} +//! @determinism{run_to_run} +//! +//! Performance +//! ==================================== +//! +//! @linear_performance{reduction, reduce-by-key, and run-length encode} +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceReduce that accept an environment (except ``ReduceByKey``) can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::ReducePolicy`, as shown in the +//! example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin reduce-policy-selector +//! :end-before: example-end reduce-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin reduce-tuning +//! :end-before: example-end reduce-tuning +//! +//! ``DeviceReduce::ReduceByKey`` can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::ReduceByKeyPolicy`, as shown in the +//! example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_by_key_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin reduce-by-key-policy-selector +//! :end-before: example-end reduce-by-key-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_by_key_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin reduce-by-key-tuning +//! :end-before: example-end reduce-by-key-tuning +//! +//! Deferred problem sizes +//! ==================================== +//! +//! ``Reduce``, ``Sum``, ``Min``, ``Max``, and ``TransformReduce`` allow specifying the problem size from a value that +//! resides in device memory through a single-value ``cuda::args::deferred`` argument. The deferred source may be a +//! device pointer, a one-element span, or a random-access fancy iterator whose element is a non-``bool`` 32- or 64-bit +//! integer. +//! +//! The problem size is read in stream order by the reduction kernels. Work that produces the count in the same stream +//! is ordered automatically; a producer in another stream requires an event or an equivalent dependency. The source +//! and its value must remain accessible and unchanged until all reduction kernels complete. The value must be +//! nonnegative and ``[d_in, d_in + num_items)`` must be accessible. A temporary-storage query does not dereference the +//! deferred source. +//! +//! Deferred reductions are CUDA Graph capturable. The pointed-to count may change between graph replays without +//! updating or recapturing the graph. Compile-time and runtime bounds are accepted as caller preconditions, but do not +//! currently change temporary storage, grid dimensions, or pass selection. Since the problem size is not available to +//! the host, the first reduction pass launches CUB's maximum grid and may launch more blocks than the actual problem +//! size needs. Extra blocks return without reducing input. +//! +//! Determinism +//! ==================================== +//! +//! ``cub::DeviceReduce`` supports all three :ref:`determinism guarantees `; the +//! default is ``run_to_run``. +//! +//! - ``run_to_run`` (the default) is reproducible because, for a given GPU, every launch with the same input, +//! build, and launch configuration selects the *same* tuning policy and therefore performs the *same* fixed +//! reduction tree: the input is partitioned into the same tiles, mapped onto the same thread blocks, and the +//! partial results are combined in the same fixed order on every run — no atomics or other run-dependent +//! ordering are involved. Because floating-point addition is only pseudo-associative, that fixed combining order +//! is what makes the result bitwise-identical from one run to the next. The combining order is tied to the +//! tuning and the partition, so it can change on a *different* GPU architecture (or under a different +//! user-provided tuning); that is exactly the cross-architecture reproducibility that ``gpu_to_gpu`` adds on top. +//! - ``gpu_to_gpu`` reproducibility depends on the type and operator: +//! +//! - ``float`` and ``double`` with ``cuda::std::plus`` use a dedicated, hardware-independent implementation +//! based on a `Reproducible Floating-point Accumulator (RFA) +//! `__: +//! input values are grouped into a fixed number of exponent-range bins and accumulated in that fixed, +//! hardware-independent order, so the same inputs yield the same bits on any GPU architecture. +//! - Exactly-associative cases — integral types with a known CUDA binary operator, and ``float``/``double`` +//! with ``min``/``max`` — are already identical across GPUs, so the request is satisfied by the (faster) +//! ``run_to_run`` path. +//! - All other type/operator combinations are rejected at compile time. +//! +//! - ``not_guaranteed`` uses an atomic accumulation kernel when the conditions for it are met: a contiguous output +//! iterator, ``cuda::std::plus``, an accumulator of at least 4 bytes, and an output type equal to the +//! accumulator type. Atomics combine partial results in whatever order the hardware schedules them, which can +//! differ between runs — hence no run-to-run guarantee, but typically the fastest option. When those conditions +//! are not met, the call falls back to ``run_to_run`` rather than failing. +//! +//! @endrst +struct DeviceReduce +{ +private: + template + CUB_RUNTIME_FUNCTION static cudaError_t reduce_impl( + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + ReductionOpT reduction_op, + TransformOpT transform_op, + T init, + ::cuda::execution::determinism::__determinism_holder_t, + const EnvT& env) + { + using args_traits_t = ::cuda::args::__traits; + using offset_t = detail::choose_offset_t; + using accum_t = decltype(detail::reduce::select_accum_t( + static_cast(nullptr))); + + if constexpr (Determinism == ::cuda::execution::determinism::__determinism_t::__gpu_to_gpu) + { + // Only instantiated with `plus`; RFA hardcodes `deterministic_sum_t`. + (void) reduction_op; + using default_policy_selector = detail::reduce:: + policy_selector_from_types, Determinism>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::rfa::dispatch( + storage, + bytes, + d_in, + d_out, + detail::make_num_items_dispatch_arg(num_items), + init, + stream, + transform_op, + policy_selector); + }); + } + else if constexpr (Determinism == ::cuda::execution::determinism::__determinism_t::__not_guaranteed) + { + using default_policy_selector = + detail::reduce::policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::reduce::dispatch( + storage, + bytes, + d_in, + THRUST_NS_QUALIFIER::unwrap_contiguous_iterator(d_out), + detail::make_num_items_dispatch_arg(num_items), + reduction_op, + init, + stream, + transform_op, + policy_selector); + }); + } + else + { + using default_policy_selector = + detail::reduce::policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::reduce::dispatch( + storage, + bytes, + d_in, + d_out, + detail::make_num_items_dispatch_arg(num_items), + reduction_op, + init, + stream, + transform_op, + policy_selector); + }); + } + } + + //! @brief Internal implementation shared by Reduce and TransformReduce env overloads + template ( + static_cast(nullptr)))> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t __transform_reduce( + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + ReductionOpT reduction_op, + TransformOpT transform_op, + T init, + const EnvT& env) + { + static_assert(!::cuda::std::execution::__queryable_with, + "Determinism should be used inside requires to have an effect."); + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using default_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + constexpr auto gpu_gpu_determinism = + ::cuda::std::is_same_v; + + // integral types are always gpu-to-gpu deterministic if reduction operator is a simple cuda binary + // operator, so fallback to run-to-run determinism + constexpr auto integral_fallback = + gpu_gpu_determinism && ::cuda::std::is_integral_v && (detail::is_cuda_binary_operator); + + // use gpu-to-gpu determinism only for float and double types with ::cuda::std::plus operator + constexpr auto float_double_plus = + gpu_gpu_determinism && detail::is_one_of_v && detail::is_cuda_std_plus_v; + + constexpr auto float_double_min_max_fallback = + gpu_gpu_determinism + && detail::is_one_of_v && detail::is_cuda_minimum_maximum_v; + + constexpr auto supported = + integral_fallback || float_double_plus || float_double_min_max_fallback || !gpu_gpu_determinism; + + // gpu_to_gpu determinism is only supported for integral types with cuda operators, or + // float and double types with ::cuda::std::plus operator + static_assert(supported, "gpu_to_gpu determinism is unsupported"); + + if constexpr (!supported) + { + return cudaErrorNotSupported; + } + else + { + constexpr auto no_determinism = detail::is_non_deterministic_v; + + // Certain conditions must be met to be able to use the non-deterministic + // kernel. The output iterator must be a contiguous iterator, the reduction + // operator must be plus (for now), and the output type must match the + // accumulator type. The non-deterministic kernel atomically accumulates + // directly into the output, so it cannot preserve AccumT accumulation + // semantics when the output object has a different type. Additionally, since + // atomics for types of size < 4B are emulated, they perform poorly, so we fall + // back to the run-to-run determinism. + using OutputT = cub::detail::non_void_value_t>; + + constexpr auto is_contiguous_fallback = + !no_determinism || THRUST_NS_QUALIFIER::is_contiguous_iterator_v; + constexpr auto is_plus_fallback = !no_determinism || detail::is_cuda_std_plus_v; + constexpr auto is_4b_or_greater = !no_determinism || sizeof(AccumT) >= 4; + constexpr auto is_output_accum = !no_determinism || ::cuda::std::is_same_v; + + // If the conditions for gpu-to-gpu determinism or non-deterministic + // reduction are not met, we fall back to run-to-run determinism. + using determinism_t = ::cuda::std::conditional_t< + (gpu_gpu_determinism && (integral_fallback || float_double_min_max_fallback)) + || (no_determinism && !(is_contiguous_fallback && is_plus_fallback && is_4b_or_greater && is_output_accum)), + ::cuda::execution::determinism::run_to_run_t, + default_determinism_t>; + + return reduce_impl(d_in, d_out, num_items, reduction_op, transform_op, init, determinism_t{}, env); + } + } + + template + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t __minmax_reduce( + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + ReductionOpT reduction_op, + InitValueT init, + const EnvT& env) + { + static_assert(!::cuda::std::execution::__queryable_with, + "Determinism should be used inside requires to have an effect."); + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + // Static assert to reject gpu_to_gpu determinism since it's not properly implemented + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported"); + + // TODO(NaderAlAwar): Relax this once non-deterministic implementation for min / max is available + using determinism_t = ::cuda::execution::determinism::run_to_run_t; + + return reduce_impl(d_in, d_out, num_items, reduction_op, ::cuda::std::identity{}, init, determinism_t{}, env); + } + +public: + //! @rst + //! Computes a device-wide reduction using the specified binary ``reduction_op`` functor and initial value ``init``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Does not support binary reduction operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a user-defined min-reduction of a + //! device vector of ``int`` data elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_out; // e.g., [-] + //! CustomMin min_op; + //! int init; // e.g., INT_MAX + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::Reduce( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, num_items, min_op, init); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run reduction + //! cub::DeviceReduce::Reduce( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, num_items, min_op, init); + //! + //! // d_out <-- [0] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam T + //! **[inferred]** Data element type that is convertible to the `value` type of `InputIteratorT` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] init + //! Initial value of the reduction + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t Reduce( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + ReductionOpT reduction_op, + T init, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::Reduce"); + + return detail::reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + detail::make_num_items_dispatch_arg(num_items), + reduction_op, + init, + stream); + } + + //! @rst + //! Computes a device-wide reduction using the specified binary ``reduction_op`` functor and initial value ``init``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Does not support binary reduction operators that are non-commutative. + //! - By default, provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! To request "gpu-to-gpu" determinism, pass ``cuda::execution::require(cuda::execution::determinism::gpu_to_gpu)`` + //! as the `env` parameter. + //! To request "not-guaranteed" determinism, pass + //! ``cuda::execution::require(cuda::execution::determinism::not_guaranteed)`` as the `env` parameter. + //! The non-deterministic implementation is only used when the output type matches the accumulator type. + //! The accumulator type is the decayed result type of invoking ``reduction_op`` with the initial value and an input + //! value. For example, reducing ``std::uint8_t`` input with ``cuda::std::plus`` and a ``std::uint8_t`` initial + //! value accumulates in ``int`` due to integer promotion. If the output type does not match the accumulator type, + //! CUB falls back to run-to-run determinism. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a user-defined min-reduction of a + //! device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin reduce-env-determinism + //! :end-before: example-end reduce-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam T + //! **[inferred]** Data element type that is convertible to the `value` type of `InputIteratorT` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] init + //! Initial value of the reduction + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t Reduce( + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + ReductionOpT reduction_op, + T init, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::Reduce"); + return __transform_reduce(d_in, d_out, num_items, reduction_op, ::cuda::std::identity{}, init, env); + } + + //! @rst + //! Computes a device-wide sum using the addition (``+``) operator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``0`` as the initial value of the reduction. + //! - Does not support ``+`` operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! To request "gpu-to-gpu" determinism, pass ``cuda::execution::require(cuda::execution::determinism::gpu_to_gpu)`` + //! as the `env` parameter. + //! To request "not-guaranteed" determinism, pass + //! ``cuda::execution::require(cuda::execution::determinism::not_guaranteed)`` as the `env` parameter. + //! The non-deterministic implementation is only used when the output type matches the accumulator type. + //! The accumulator type is the decayed result type of adding the implicit initial value, whose type is the output + //! type, to an input value. For example, summing ``std::uint8_t`` input into ``std::uint8_t`` output accumulates in + //! ``int`` due to integer promotion. + //! If the output type does not match the accumulator type, CUB falls back to run-to-run determinism. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a user-defined min-reduction of a + //! device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sum-env-determinism + //! :end-before: example-end sum-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is `cuda::std::execution::env<>`. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Sum(InputIteratorT d_in, OutputIteratorT d_out, NumItemsT num_items, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::Sum"); + using OutputT = cub::detail::non_void_value_t>; + + return __transform_reduce(d_in, d_out, num_items, ::cuda::std::plus<>{}, ::cuda::std::identity{}, OutputT{}, env); + } + + //! @rst + //! Computes a device-wide sum using the addition (``+``) operator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``0`` as the initial value of the reduction. + //! - Does not support ``+`` operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the sum-reduction of a device vector + //! of ``int`` data elements. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_out; // e.g., [-] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::Sum( + //! d_temp_storage, temp_storage_bytes, d_in, d_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sum-reduction + //! cub::DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items); + //! + //! // d_out <-- [38] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t + Sum(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::Sum"); + + // The output value type + using OutputT = cub::detail::non_void_value_t>; + + using init_value_t = OutputT; + + return detail::reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + detail::make_num_items_dispatch_arg(num_items), + ::cuda::std::plus<>{}, + init_value_t{}, // zero-initialize + stream); + } + + //! @rst + //! Computes a device-wide minimum using the less-than (``<``) operator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``cuda::std::numeric_limits::max()`` as the initial value of the reduction. + //! - Does not support ``<`` operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the min-reduction of a device vector of ``int`` data elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_out; // e.g., [-] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::Min( + //! d_temp_storage, temp_storage_bytes, d_in, d_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run min-reduction + //! cub::DeviceReduce::Min( + //! d_temp_storage, temp_storage_bytes, d_in, d_out, num_items); + //! + //! // d_out <-- [0] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t + Min(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::Min"); + + using InputT = detail::it_value_t; + using init_value_t = InputT; + using limits_t = ::cuda::std::numeric_limits; +#ifndef CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX + static_assert(limits_t::is_specialized, + "cub::DeviceReduce::Min uses cuda::std::numeric_limits::max() as initial " + "value, but cuda::std::numeric_limits is not specialized for the iterator's value type. This is " + "probably a bug and you should specialize cuda::std::numeric_limits. Define " + "CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX to suppress this check."); +#endif // CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX + + return detail::reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + detail::make_num_items_dispatch_arg(num_items), + ::cuda::minimum<>{}, + limits_t::max(), + stream); + } + + //! @rst + //! Computes a device-wide minimum using the less-than (``<``) operator. The result is written to the output + //! iterator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``cuda::std::numeric_limits::max()`` as the initial value of the reduction. + //! - Provides determinism based on the environment's determinism requirements. + //! To request "run-to-run" determinism, pass ``cuda::execution::require(cuda::execution::determinism::run_to_run)`` + //! as the `env` parameter. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the min-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin min-env-determinism + //! :end-before: example-end min-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is `cuda::std::execution::env<>`. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is `cuda::std::execution::env{}`. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Min(InputIteratorT d_in, OutputIteratorT d_out, NumItemsT num_items, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::Min"); + using OutputT = cub::detail::non_void_value_t>; + using limits_t = ::cuda::std::numeric_limits; +#ifndef CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX + static_assert(limits_t::is_specialized, + "cub::DeviceReduce::Min uses cuda::std::numeric_limits::max() as initial " + "value, but cuda::std::numeric_limits is not specialized for the iterator's value type. This is " + "probably a bug and you should specialize cuda::std::numeric_limits. Define " + "CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX to suppress this check."); +#endif // CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX + return __minmax_reduce(d_in, d_out, num_items, ::cuda::minimum<>{}, limits_t::max(), env); + } + +private: + template + [[nodiscard]] _CCCL_API static _CCCL_CONSTEVAL bool __validate_determinism_streaming_reduce() noexcept + { + static_assert(!::cuda::std::execution::__queryable_with, + "Determinism should be used inside requires to have an effect."); + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + // Reject gpu_to_gpu determinism since it's not properly implemented + return !::cuda::std::is_same_v; + } + + template + CUB_RUNTIME_FUNCTION static cudaError_t __arg_min( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + ExtremumOutIteratorT d_min_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + CompareOpT compare_op, + const EnvT& env) + { + static_assert(__validate_determinism_streaming_reduce(), "gpu_to_gpu determinism is not supported"); + + using PerPartitionOffsetT = int; // used by the kernel to index within one partition + using GlobalOffsetT = ::cuda::std::int64_t; // in the range [d_in, d_in + num_items) + using reduce_op_t = detail::arg_reduce_op; + + return detail::dispatch_with_env( + d_temp_storage, temp_storage_bytes, env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::reduce::dispatch_streaming_arg_reduce( + storage, + bytes, + d_in, + d_min_out, + d_index_out, + static_cast(num_items), + reduce_op_t{compare_op}, + stream, + tuning_env); + }); + } + + template + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t __arg_min( + InputIteratorT d_in, + ExtremumOutIteratorT d_min_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + CompareOpT compare_op, + const EnvT& env) + { + static_assert(__validate_determinism_streaming_reduce(), "gpu_to_gpu determinism is not supported"); + + using PerPartitionOffsetT = int; // used by the kernel to index within one partition + using GlobalOffsetT = ::cuda::std::int64_t; // in the range [d_in, d_in + num_items) + using reduce_op_t = detail::arg_reduce_op; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::reduce::dispatch_streaming_arg_reduce( + storage, + bytes, + d_in, + d_min_out, + d_index_out, + static_cast(num_items), + reduce_op_t{compare_op}, + stream, + tuning_env); + }); + } + +public: + //! @rst + //! Finds the first device-wide minimum based on a given comparison operator and also returns the index of that item. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The minimum is written to ``d_min_out`` + //! - The offset of the returned item is written to ``d_index_out``, the offset type being written is of type + //! ``cuda::std::int64_t``. + //! - For zero-length inputs, the index ``1`` is written to ``d_index_out`` and, if ``compare_op`` is + //! ``cuda::std::less`` and ``cuda::std::numeric_limits::is_specialized is ``true``, + //! ``cuda::std::numeric_limits::max()`` is written to ``d_min_out``, otherwise ``T{}``. + //! - Does not support comparison operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_min_out`` nor ``d_index_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmin-reduction of a device vector + //! of ``int`` data elements. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, -7, 5, 3, 1, -9] + //! int *d_min_out; // memory for the minimum value + //! cuda::std::int64_t *d_index_out; // memory for the index of the returned value + //! ... + //! + //! // Define the comparison operator + //! struct abs_less_t { + //! template + //! __host__ __device__ bool operator()(const T& a, const T& b) const { + //! return cuda::std::abs(a) < cuda::std::abs(b); + //! } + //! }; + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_min_out, d_index_out, + //! num_items, abs_less_t{}); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run argmin-reduction + //! cub::DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_min_out, d_index_out, + //! num_items, abs_less_t{}); + //! + //! // d_min_out <-- 1 + //! // d_index_out <-- 5 + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items + //! (of some type `T`) @iterator + //! + //! @tparam ExtremumOutIteratorT + //! **[inferred]** Output iterator type for recording minimum value + //! + //! @tparam IndexOutIteratorT + //! **[inferred]** Output iterator type for recording index of the returned value + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @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_in + //! Iterator to the input sequence of data items + //! + //! @param[out] d_min_out + //! Iterator to which the minimum value is written + //! + //! @param[out] d_index_out + //! Iterator to which the index of the returned value is written + //! + //! @param[in] compare_op + //! Comparison operator returning ``true`` if the first argument is less than the second + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + // TODO(bgruber): this constraint is not accurate, since the implementation will compare the value types of + // ExtremumOutIteratorT, which is wrong IMO + _CCCL_TEMPLATE(typename InputIteratorT, + typename ExtremumOutIteratorT, + typename IndexOutIteratorT, + typename CompareOpT, + typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES((::cuda::std::indirectly_comparable) ) + CUB_RUNTIME_FUNCTION static cudaError_t ArgMin( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + ExtremumOutIteratorT d_min_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + CompareOpT compare_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::ArgMin"); + return __arg_min(d_temp_storage, temp_storage_bytes, d_in, d_min_out, d_index_out, num_items, compare_op, env); + } + + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @endrst + //! + //! @overload + //! @note Uses ``cuda::std::less`` as comparison operator + _CCCL_TEMPLATE(typename InputIteratorT, + typename ExtremumOutIteratorT, + typename IndexOutIteratorT, + typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES((!::cuda::std::indirectly_comparable) ) + CUB_RUNTIME_FUNCTION static cudaError_t ArgMin( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + ExtremumOutIteratorT d_min_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + return ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_min_out, d_index_out, num_items, ::cuda::std::less{}, env); + } + +public: + //! @rst + //! Finds the first device-wide minimum using the less-than (``<``) operator and also returns the index of that item. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The minimum is written to ``d_min_out`` + //! - The offset of the returned item is written to ``d_index_out``, the offset type being written is of type + //! ``cuda::std::int64_t``. + //! - For zero-length inputs, the index ``1`` is written to ``d_index_out`` and, if ``compare_op`` is + //! ``cuda::std::less`` and ``cuda::std::numeric_limits::is_specialized is ``true``, + //! ``cuda::std::numeric_limits::max()`` is written to ``d_min_out``, otherwise ``T{}``. + //! - Does not support ``<`` operators that are non-commutative. + //! - Provides determinism based on the environment's determinism requirements. + //! To request "run-to-run" determinism, pass ``cuda::execution::require(cuda::execution::determinism::run_to_run)`` + //! as the `env` parameter. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_min_out`` nor ``d_index_out``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmin-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin argmin-env-determinism + //! :end-before: example-end argmin-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items + //! (of some type `T`) @iterator + //! + //! @tparam ExtremumOutIteratorT + //! **[inferred]** Output iterator type for recording minimum value + //! + //! @tparam IndexOutIteratorT + //! **[inferred]** Output iterator type for recording index of the returned value + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Iterator to the input sequence of data items + //! + //! @param[out] d_min_out + //! Iterator to which the minimum value is written + //! + //! @param[out] d_index_out + //! Iterator to which the index of the returned value is written + //! + //! @param[in] compare_op + //! Comparison operator returning ``true`` if the first argument is less than the second + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + // TODO(bgruber): this constraint is not accurate, since the implementation will compare the value types of + // ExtremumOutIteratorT, which is wrong IMO + _CCCL_TEMPLATE(typename InputIteratorT, + typename ExtremumOutIteratorT, + typename IndexOutIteratorT, + typename CompareOpT, + typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES((::cuda::std::indirectly_comparable) ) + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ArgMin( + InputIteratorT d_in, + ExtremumOutIteratorT d_min_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + CompareOpT compare_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::ArgMin"); + return __arg_min(d_in, d_min_out, d_index_out, num_items, compare_op, env); + } + + //! @overload + //! @note Uses ``cuda::std::less`` as comparison operator + // TODO(bgruber): this constraint is not accurate, since the implementation will compare the value types of + // ExtremumOutIteratorT, which is wrong IMO + _CCCL_TEMPLATE(typename InputIteratorT, + typename ExtremumOutIteratorT, + typename IndexOutIteratorT, + typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES((!::cuda::std::indirectly_comparable) ) + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ArgMin( + InputIteratorT d_in, + ExtremumOutIteratorT d_min_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::ArgMin"); + return __arg_min(d_in, d_min_out, d_index_out, num_items, ::cuda::std::less{}, env); + } + + //! @rst + //! Finds the first device-wide minimum using the less-than (``<``) operator, also returning the index of that item. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The output value type of ``d_out`` is ``cub::KeyValuePair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The minimum is written to ``d_out.value`` and its offset in the input array is written to ``d_out.key``. + //! - The ``{1, cuda::std::numeric_limits::max()}`` tuple is produced for zero-length inputs + //! + //! - Does not support ``<`` operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmin-reduction of a device vector + //! of ``int`` data elements. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! KeyValuePair *d_argmin; // e.g., [{-,-}] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_argmin, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run argmin-reduction + //! cub::DeviceReduce::ArgMin(d_temp_storage, temp_storage_bytes, d_in, d_argmin, num_items); + //! + //! // d_argmin <-- [{5, 0}] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items + //! (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type ``cub::KeyValuePair``) @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CCCL_DEPRECATED_BECAUSE("CUB has superseded this interface in favor of the ArgMin interface that takes two separate " + "iterators: one iterator to which the extremum is written and another iterator to which the " + "index of the found extremum is written. ") CUB_RUNTIME_FUNCTION static cudaError_t + ArgMin(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + int num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::ArgMin"); + + // Signed integer type for global offsets + using OffsetT = int; + + // The input type + using InputValueT = cub::detail::it_value_t; + + // The output tuple type + using OutputTupleT = cub::detail::non_void_value_t>; + + using AccumT = OutputTupleT; + + using init_value_t = detail::reduce::empty_problem_init_t; + + // The output value type + using OutputValueT = typename OutputTupleT::Value; + + // Wrapped input iterator to produce index-value tuples + using ArgIndexInputIteratorT = ArgIndexInputIterator; + + ArgIndexInputIteratorT d_indexed_in(d_in); + + // Initial value + init_value_t initial_value{AccumT(1, ::cuda::std::numeric_limits::max())}; + + return detail::reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_indexed_in, + d_out, + OffsetT{num_items}, + detail::arg_min{}, + initial_value, + stream); + } + + //! @rst + //! Computes a device-wide maximum using the greater-than (``>``) operator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``cuda::std::numeric_limits::lowest()`` as the initial value of the reduction. + //! - Does not support ``>`` operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the max-reduction of a device vector of ``int`` data elements. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_max; // e.g., [-] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::Max(d_temp_storage, temp_storage_bytes, d_in, d_max, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run max-reduction + //! cub::DeviceReduce::Max(d_temp_storage, temp_storage_bytes, d_in, d_max, num_items); + //! + //! // d_max <-- [9] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t + Max(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::Max"); + + using InputT = detail::it_value_t; + using init_value_t = InputT; + using limits_t = ::cuda::std::numeric_limits; +#ifndef CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX + static_assert(limits_t::is_specialized, + "cub::DeviceReduce::Max uses cuda::std::numeric_limits::lowest() as " + "initial value, but cuda::std::numeric_limits is not specialized for the iterator's value type. This " + "is probably a bug and you should specialize cuda::std::numeric_limits. Define " + "CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX to suppress this check."); +#endif // CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX + + return detail::reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + detail::make_num_items_dispatch_arg(num_items), + ::cuda::maximum<>{}, + limits_t::lowest(), + stream); + } + + //! @rst + //! Computes a device-wide maximum using the greater-than (``>``) operator. The result is written to the output + //! iterator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``cuda::std::numeric_limits::lowest()`` as the initial value of the reduction. + //! - Provides determinism based on the environment's determinism requirements. + //! To request "run-to-run" determinism, pass ``cuda::execution::require(cuda::execution::determinism::run_to_run)`` + //! as the `env` parameter. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the max-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin max-env-determinism + //! :end-before: example-end max-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is `cuda::std::execution::env<>`. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Max(InputIteratorT d_in, OutputIteratorT d_out, NumItemsT num_items, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::Max"); + using OutputT = cub::detail::non_void_value_t>; + using limits_t = ::cuda::std::numeric_limits; +#ifndef CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX + static_assert( + limits_t::is_specialized, + "cub::DeviceReduce::Max uses cuda::std::numeric_limits::lowest() as initial value, " + "but cuda::std::numeric_limits is not specialized for the iterator's value type. This is probably a bug and you " + "should specialize cuda::std::numeric_limits. Define " + "CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX to suppress this check."); +#endif // CCCL_SUPPRESS_NUMERIC_LIMITS_CHECK_IN_CUB_DEVICE_REDUCE_MIN_MAX + return __minmax_reduce(d_in, d_out, num_items, ::cuda::maximum<>{}, limits_t::lowest(), env); + } + + //! @rst + //! Finds the first device-wide maximum based on a given comparison operator and also returns the index of that item. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The maximum is written to ``d_max_out`` + //! - The offset of the returned item is written to ``d_index_out``, the offset type being written is of type + //! ``cuda::std::int64_t``. + //! - For zero-length inputs, the index ``1`` is written to ``d_index_out`` and, if ``compare_op`` is + //! ``cuda::std::less`` and ``cuda::std::numeric_limits::is_specialized is ``true``, + //! ``cuda::std::numeric_limits::lowest()`` is written to ``d_min_out``, otherwise ``T{}``. + //! - Does not support ``>`` operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmax-reduction of a device vector + //! of `int` data elements. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, -7, 5, 3, 1, -9] + //! int *d_max_out; // memory for the maximum value + //! cuda::std::int64_t *d_index_out; // memory for the index of the returned value + //! ... + //! + //! // Define the comparison operator + //! struct abs_less_t { + //! template + //! __host__ __device__ bool operator()(const T& a, const T& b) const { + //! return cuda::std::abs(a) < cuda::std::abs(b); + //! } + //! }; + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::ArgMax( + //! d_temp_storage, temp_storage_bytes, d_in, d_max_out, d_index_out, num_items, abs_less_t{}); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run argmax-reduction + //! cub::DeviceReduce::ArgMax( + //! d_temp_storage, temp_storage_bytes, d_in, d_max_out, d_index_out, num_items, abs_less_t{}); + //! + //! // d_max_out <-- -9 + //! // d_index_out <-- 6 + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items (of some type `T`) @iterator + //! + //! @tparam ExtremumOutIteratorT + //! **[inferred]** Output iterator type for recording maximum value + //! + //! @tparam IndexOutIteratorT + //! **[inferred]** Output iterator type for recording index of the returned value + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_max_out + //! Iterator to which the maximum value is written + //! + //! @param[out] d_index_out + //! Iterator to which the index of the returned value is written + //! + //! @param[in] compare_op + //! Comparison operator returning ``true`` if the first argument is less than the second + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + // TODO(bgruber): this constraint is not accurate, since the implementation will compare the value types of + // ExtremumOutIteratorT, which is wrong IMO + _CCCL_TEMPLATE(typename InputIteratorT, + typename ExtremumOutIteratorT, + typename IndexOutIteratorT, + typename CompareOpT, + typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES((::cuda::std::indirectly_comparable) ) + CUB_RUNTIME_FUNCTION static cudaError_t ArgMax( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + ExtremumOutIteratorT d_max_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + CompareOpT compare_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::ArgMax"); + return __arg_min( + d_temp_storage, temp_storage_bytes, d_in, d_max_out, d_index_out, num_items, detail::swap_args{compare_op}, env); + } + + //! @rst + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! @overload + //! @note Uses ``cuda::std::less`` as comparison operator + //! @endrst + _CCCL_TEMPLATE(typename InputIteratorT, + typename ExtremumOutIteratorT, + typename IndexOutIteratorT, + typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES((!::cuda::std::indirectly_comparable) ) + CUB_RUNTIME_FUNCTION static cudaError_t ArgMax( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + ExtremumOutIteratorT d_max_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::ArgMax"); + return __arg_min( + d_temp_storage, temp_storage_bytes, d_in, d_max_out, d_index_out, num_items, ::cuda::std::greater{}, env); + } + + //! @rst + //! Finds the first device-wide maximum using the greater-than (``>``) + //! operator, also returning the index of that item + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The output value type of ``d_out`` is ``cub::KeyValuePair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The maximum is written to ``d_out.value`` and its offset in the input + //! array is written to ``d_out.key``. + //! - The ``{1, cuda::std::numeric_limits::lowest()}`` tuple is produced for zero-length inputs + //! + //! - Does not support ``>`` operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmax-reduction of a device vector + //! of `int` data elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! KeyValuePair *d_argmax; // e.g., [{-,-}] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::ArgMax( + //! d_temp_storage, temp_storage_bytes, d_in, d_argmax, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run argmax-reduction + //! cub::DeviceReduce::ArgMax( + //! d_temp_storage, temp_storage_bytes, d_in, d_argmax, num_items); + //! + //! // d_argmax <-- [{6, 9}] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `cub::KeyValuePair`) @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CCCL_DEPRECATED_BECAUSE("CUB has superseded this interface in favor of the ArgMax interface that takes two separate " + "iterators: one iterator to which the extremum is written and another iterator to which the " + "index of the found extremum is written. ") CUB_RUNTIME_FUNCTION static cudaError_t + ArgMax(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + int num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::ArgMax"); + + // Signed integer type for global offsets + using OffsetT = int; + + // The input type + using InputValueT = cub::detail::it_value_t; + + // The output tuple type + using OutputTupleT = cub::detail::non_void_value_t>; + + using AccumT = OutputTupleT; + + // The output value type + using OutputValueT = typename OutputTupleT::Value; + + using init_value_t = detail::reduce::empty_problem_init_t; + + // Wrapped input iterator to produce index-value tuples + using ArgIndexInputIteratorT = ArgIndexInputIterator; + + ArgIndexInputIteratorT d_indexed_in(d_in); + + // Initial value + init_value_t initial_value{AccumT(1, ::cuda::std::numeric_limits::lowest())}; + + return detail::reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_indexed_in, + d_out, + OffsetT{num_items}, + detail::arg_max{}, + initial_value, + stream); + } + + //! @rst + //! Finds the first device-wide maximum using the greater-than (``>``) operator and also returns the index of that + //! item. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The maximum is written to ``d_max_out`` + //! - The offset of the returned item is written to ``d_index_out``, the offset type being written is of type + //! ``cuda::std::int64_t``. + //! - For zero-length inputs, the index ``1`` is written to ``d_index_out`` and, if ``compare_op`` is + //! ``cuda::std::less`` and ``cuda::std::numeric_limits::is_specialized is ``true``, + //! ``cuda::std::numeric_limits::lowest()`` is written to ``d_min_out``, otherwise ``T{}``. + //! - Does not support ``>`` operators that are non-commutative. + //! - Provides determinism based on the environment's determinism requirements. + //! To request "run-to-run" determinism, pass ``cuda::execution::require(cuda::execution::determinism::run_to_run)`` + //! as the `env` parameter. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_max_out`` nor ``d_index_out``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmax-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin argmax-env-determinism + //! :end-before: example-end argmax-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items + //! (of some type `T`) @iterator + //! + //! @tparam ExtremumOutIteratorT + //! **[inferred]** Output iterator type for recording maximum value + //! + //! @tparam IndexOutIteratorT + //! **[inferred]** Output iterator type for recording index of the returned value + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Iterator to the input sequence of data items + //! + //! @param[out] d_max_out + //! Iterator to which the maximum value is written + //! + //! @param[out] d_index_out + //! Iterator to which the index of the returned value is written + //! + //! @param[in] compare_op + //! Comparison operator returning ``true`` if the first argument is less than the second + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + // TODO(bgruber): this constraint is not accurate, since the implementation will compare the value types of + // ExtremumOutIteratorT, which is wrong IMO + _CCCL_TEMPLATE(typename InputIteratorT, + typename ExtremumOutIteratorT, + typename IndexOutIteratorT, + typename CompareOpT, + typename EnvT = ::cuda::std::execution::env<>) + _CCCL_REQUIRES((::cuda::std::indirectly_comparable) ) + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ArgMax( + InputIteratorT d_in, + ExtremumOutIteratorT d_max_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + CompareOpT compare_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::ArgMax"); + return __arg_min(d_in, d_max_out, d_index_out, num_items, detail::swap_args{compare_op}, env); + } + + //! @overload + //! @note Uses ``cuda::std::less`` as comparison operator + template , + // TODO(bgruber): this constraint is not accurate, since the implementation will compare the value types of + // ExtremumOutIteratorT, which is wrong IMO + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + ArgMax(InputIteratorT d_in, + ExtremumOutIteratorT d_max_out, + IndexOutIteratorT d_index_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::ArgMax"); + return __arg_min(d_in, d_max_out, d_index_out, num_items, ::cuda::std::greater{}, env); + } + + //! @rst + //! Fuses transform and reduce operations + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Does not support binary reduction operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a user-defined min-reduction of a + //! device vector of `int` data elements. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! thrust::device_vector in = { 1, 2, 3, 4 }; + //! thrust::device_vector out(1); + //! + //! size_t temp_storage_bytes = 0; + //! uint8_t *d_temp_storage = nullptr; + //! + //! const int init = 42; + //! + //! cub::DeviceReduce::TransformReduce( + //! d_temp_storage, + //! temp_storage_bytes, + //! in.begin(), + //! out.begin(), + //! in.size(), + //! cuda::std::plus<>{}, + //! square_t{}, + //! init); + //! + //! thrust::device_vector temp_storage(temp_storage_bytes); + //! d_temp_storage = temp_storage.data().get(); + //! + //! cub::DeviceReduce::TransformReduce( + //! d_temp_storage, + //! temp_storage_bytes, + //! in.begin(), + //! out.begin(), + //! in.size(), + //! cuda::std::plus<>{}, + //! square_t{}, + //! init); + //! + //! // out[0] <-- 72 + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam TransformOpT + //! **[inferred]** Unary reduction functor type having member `auto operator()(const T &a)` + //! + //! @tparam T + //! **[inferred]** Data element type that is convertible to the `value` type of `InputIteratorT` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] transform_op + //! Unary transform functor + //! + //! @param[in] init + //! Initial value of the reduction + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t TransformReduce( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + ReductionOpT reduction_op, + TransformOpT transform_op, + T init, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::TransformReduce"); + + return detail::reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + detail::make_num_items_dispatch_arg(num_items), + reduction_op, + init, + stream, + transform_op); + } + + //! @rst + //! Computes a device-wide reduction using the specified binary ``reduction_op`` functor, + //! a unary ``transform_op`` functor, and initial value ``init``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Does not support binary reduction operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! To request "gpu-to-gpu" determinism, pass ``cuda::execution::require(cuda::execution::determinism::gpu_to_gpu)`` + //! as the `env` parameter. + //! To request "not-guaranteed" determinism, pass + //! ``cuda::execution::require(cuda::execution::determinism::not_guaranteed)`` as the `env` parameter. + //! The non-deterministic implementation is only used when the output type matches the accumulator type. + //! The accumulator type is the decayed result type of invoking ``reduction_op`` with the initial value and the + //! transformed input value. For example, reducing transformed ``std::uint8_t`` values with ``cuda::std::plus`` and + //! a + //! ``std::uint8_t`` initial value accumulates in ``int`` due to integer promotion. + //! If the output type does not match the accumulator type, CUB falls back to run-to-run determinism. + //! - The range ``[d_in, d_in + num_items)`` shall not overlap ``d_out``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a transform-reduction with determinism + //! of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin transform-reduce-env-determinism + //! :end-before: example-end transform-reduce-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam TransformOpT + //! **[inferred]** Unary reduction functor type having member `auto operator()(const T &a)` + //! + //! @tparam T + //! **[inferred]** Data element type that is convertible to the `value` type of `InputIteratorT` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of ``d_in``) + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] transform_op + //! Unary transform functor + //! + //! @param[in] init + //! Initial value of the reduction + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t TransformReduce( + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + ReductionOpT reduction_op, + TransformOpT transform_op, + T init, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::TransformReduce"); + return __transform_reduce(d_in, d_out, num_items, reduction_op, transform_op, init, env); + } + + //! @rst + //! Reduces segments of values, where segments are demarcated by corresponding runs of identical keys. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! This operation computes segmented reductions within ``d_values_in`` using the specified binary ``reduction_op`` + //! functor. The segments are identified by "runs" of corresponding keys in `d_keys_in`, where runs are maximal + //! ranges of consecutive, identical keys. For the *i*\ :sup:`th` run encountered, the last key of the run and + //! the corresponding value aggregate of that run are written to ``d_unique_out[i]`` and ``d_aggregates_out[i]``, + //! respectively. The total number of runs encountered is written to ``d_num_runs_out``. + //! + //! - The ``==`` equality operator is used to determine whether keys are equivalent + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! - Let ``out`` be any of + //! ``[d_unique_out, d_unique_out + *d_num_runs_out)`` + //! ``[d_aggregates_out, d_aggregates_out + *d_num_runs_out)`` + //! ``d_num_runs_out``. The ranges represented by ``out`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_values_in, d_values_in + num_items)`` nor ``out`` in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a reduce-by-key operation + //! of a device vector of ``int`` keys and values. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin reduce-by-key-env + //! :end-before: example-end reduce-by-key-env + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input keys @iterator + //! + //! @tparam UniqueOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing unique output keys @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input values @iterator + //! + //! @tparam AggregatesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output value aggregates @iterator + //! + //! @tparam NumRunsOutputIteratorT + //! **[inferred]** Output iterator type for recording the number of runs encountered @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! Supports customization of stream via ``cuda::get_stream``. + //! + //! @param[in] d_keys_in + //! Pointer to the input sequence of keys + //! + //! @param[out] d_unique_out + //! Pointer to the output sequence of unique keys (one key per run) + //! + //! @param[in] d_values_in + //! Pointer to the input sequence of corresponding values + //! + //! @param[out] d_aggregates_out + //! Pointer to the output sequence of value aggregates + //! (one aggregate per run) + //! + //! @param[out] d_num_runs_out + //! Pointer to total number of runs encountered + //! (i.e., the length of ``d_unique_out``) + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] num_items + //! Total number of associated key+value pairs + //! (i.e., the length of ``d_in_keys`` and ``d_in_values``) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ReduceByKey( + KeysInputIteratorT d_keys_in, + UniqueOutputIteratorT d_unique_out, + ValuesInputIteratorT d_values_in, + AggregatesOutputIteratorT d_aggregates_out, + NumRunsOutputIteratorT d_num_runs_out, + ReductionOpT reduction_op, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceReduce::ReduceByKey"); + + using OffsetT = detail::choose_offset_t; + using EqualityOp = ::cuda::std::equal_to<>; + using default_policy_selector = detail::reduce_by_key::policy_selector_from_types< + ReductionOpT, + ::cuda::std::__accumulator_t>, + detail::non_void_value_t>>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, cudaStream_t stream) { + return detail::reduce_by_key::dispatch( + storage, + bytes, + d_keys_in, + d_unique_out, + d_values_in, + d_aggregates_out, + d_num_runs_out, + EqualityOp{}, + reduction_op, + static_cast(num_items), + stream, + policy_selector); + }); + } + + //! @rst + //! Reduces segments of values, where segments are demarcated by corresponding runs of identical keys. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! This operation computes segmented reductions within ``d_values_in`` using the specified binary ``reduction_op`` + //! functor. The segments are identified by "runs" of corresponding keys in `d_keys_in`, where runs are maximal + //! ranges of consecutive, identical keys. For the *i*\ :sup:`th` run encountered, the last key of the run and + //! the corresponding value aggregate of that run are written to ``d_unique_out[i]`` and ``d_aggregates_out[i]``, + //! respectively. The total number of runs encountered is written to ``d_num_runs_out``. + //! + //! - The ``==`` equality operator is used to determine whether keys are equivalent + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - Let ``out`` be any of + //! ``[d_unique_out, d_unique_out + *d_num_runs_out)`` + //! ``[d_aggregates_out, d_aggregates_out + *d_num_runs_out)`` + //! ``d_num_runs_out``. The ranges represented by ``out`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_values_in, d_values_in + num_items)`` nor ``out`` in any way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the segmented reduction of ``int`` values grouped by runs of + //! associated ``int`` keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 8 + //! int *d_keys_in; // e.g., [0, 2, 2, 9, 5, 5, 5, 8] + //! int *d_values_in; // e.g., [0, 7, 1, 6, 2, 5, 3, 4] + //! int *d_unique_out; // e.g., [-, -, -, -, -, -, -, -] + //! int *d_aggregates_out; // e.g., [-, -, -, -, -, -, -, -] + //! int *d_num_runs_out; // e.g., [-] + //! CustomMin reduction_op; + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceReduce::ReduceByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_unique_out, d_values_in, + //! d_aggregates_out, d_num_runs_out, reduction_op, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run reduce-by-key + //! cub::DeviceReduce::ReduceByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_unique_out, d_values_in, + //! d_aggregates_out, d_num_runs_out, reduction_op, num_items); + //! + //! // d_unique_out <-- [0, 2, 9, 5, 8] + //! // d_aggregates_out <-- [0, 1, 6, 2, 4] + //! // d_num_runs_out <-- [5] + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input keys @iterator + //! + //! @tparam UniqueOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing unique output keys @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input values @iterator + //! + //! @tparam AggregatesOutputIterator + //! **[inferred]** Random-access output iterator type for writing output value aggregates @iterator + //! + //! @tparam NumRunsOutputIteratorT + //! **[inferred]** Output iterator type for recording the number of runs encountered @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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 + //! Pointer to the input sequence of keys + //! + //! @param[out] d_unique_out + //! Pointer to the output sequence of unique keys (one key per run) + //! + //! @param[in] d_values_in + //! Pointer to the input sequence of corresponding values + //! + //! @param[out] d_aggregates_out + //! Pointer to the output sequence of value aggregates + //! (one aggregate per run) + //! + //! @param[out] d_num_runs_out + //! Pointer to total number of runs encountered + //! (i.e., the length of ``d_unique_out``) + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] num_items + //! Total number of associated key+value pairs + //! (i.e., the length of ``d_in_keys`` and ``d_in_values``) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t ReduceByKey( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + UniqueOutputIteratorT d_unique_out, + ValuesInputIteratorT d_values_in, + AggregatesOutputIteratorT d_aggregates_out, + NumRunsOutputIteratorT d_num_runs_out, + ReductionOpT reduction_op, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceReduce::ReduceByKey"); + + using OffsetT = detail::choose_offset_t; + using EqualityOp = ::cuda::std::equal_to<>; + + return detail::reduce_by_key::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_unique_out, + d_values_in, + d_aggregates_out, + d_num_runs_out, + EqualityOp{}, + reduction_op, + static_cast(num_items), + stream); + } +}; +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_run_length_encode.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_run_length_encode.cuh new file mode 100644 index 00000000..4a1bf308 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_run_length_encode.cuh @@ -0,0 +1,620 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceRunLengthEncode provides device-wide, parallel operations for computing a run-length encoding across a +//! sequence of data items residing within device-accessible memory. + +#pragma once + +#include + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceRunLengthEncode provides device-wide, parallel operations for +//! demarcating "runs" of same-valued items within a sequence residing +//! within device-accessible memory. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! A `run-length encoding `_ +//! computes a simple compressed representation of a sequence of input elements +//! such that each maximal "run" of consecutive same-valued data items is +//! encoded as a single data value along with a count of the elements in that +//! run. +//! +//! Usage Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @cdp_class{DeviceRunLengthEncode} +//! +//! Performance +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @linear_performance{run-length encode} +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The ``Encode`` algorithm that accepts an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns an :cpp:struct:`cub::RleEncodePolicy`. +//! +//! The ``NonTrivialRuns`` algorithm that accepts an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns an :cpp:struct:`cub::RleNonTrivialRunsPolicy`, as shown +//! in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_run_length_encode_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin non-trivial-runs-policy-selector +//! :end-before: example-end non-trivial-runs-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_run_length_encode_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin non-trivial-runs-tuning +//! :end-before: example-end non-trivial-runs-tuning +//! +//! @endrst +struct DeviceRunLengthEncode +{ +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // DeviceRunLengthEncode::Encode dispatches to ReduceByKey, but we want to have a dedicated tuning policy, so we need + // to adapt the policy selector to convert the tuning policy + template +# if _CCCL_HAS_CONCEPTS() + requires detail::rle::encode::rle_encode_policy_selector +# endif // _CCCL_HAS_CONCEPTS() + struct __policy_selector_adapter + { + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const + -> ReduceByKeyPolicy + { + const RleEncodePolicy policy = PolicySelector{}(cc); + return ReduceByKeyPolicy{ + ReduceByKeyAlgorithm::lookback, + {policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + policy.lookback.lookback_delay}}; + } + }; +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Computes a run-length encoding of the sequence ``d_in``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - For the *i*\ :sup:`th` run encountered, the first key of the run and + //! its length are written to ``d_unique_out[i]`` and ``d_counts_out[i]``, respectively. + //! - The total number of runs encountered is written to ``d_num_runs_out``. + //! - The ``==`` equality operator is used to determine whether values are equivalent + //! - In-place operations are not supported. There must be no overlap between any of the provided ranges: + //! + //! - ``[d_unique_out, d_unique_out + *d_num_runs_out)`` + //! - ``[d_counts_out, d_counts_out + *d_num_runs_out)`` + //! - ``[d_num_runs_out, d_num_runs_out + 1)`` + //! - ``[d_in, d_in + num_items)`` + //! + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the run-length encoding of a sequence of ``int`` values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [0, 2, 2, 9, 5, 5, 5, 8] + //! int *d_unique_out; // e.g., [ , , , , , , , ] + //! int *d_counts_out; // e.g., [ , , , , , , , ] + //! int *d_num_runs_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRunLengthEncode::Encode( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_unique_out, d_counts_out, d_num_runs_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run encoding + //! cub::DeviceRunLengthEncode::Encode( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_unique_out, d_counts_out, d_num_runs_out, num_items); + //! + //! // d_unique_out <-- [0, 2, 9, 5, 8] + //! // d_counts_out <-- [1, 2, 1, 3, 1] + //! // d_num_runs_out <-- [5] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam UniqueOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing unique output items @iterator + //! + //! @tparam LengthsOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output counts @iterator + //! + //! @tparam NumRunsOutputIteratorT + //! **[inferred]** Output iterator type for recording the number of runs encountered @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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_in + //! Pointer to the input sequence of keys + //! + //! @param[out] d_unique_out + //! Pointer to the output sequence of unique keys (one key per run) + //! + //! @param[out] d_counts_out + //! Pointer to the output sequence of run-lengths (one count per run) + //! + //! @param[out] d_num_runs_out + //! Pointer to total number of runs + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Encode( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + UniqueOutputIteratorT d_unique_out, + LengthsOutputIteratorT d_counts_out, + NumRunsOutputIteratorT d_num_runs_out, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceRunLengthEncode::Encode"); + + using equality_op = ::cuda::std::equal_to<>; // Default == operator + using reduction_op = ::cuda::std::plus<>; // Value reduction operator + + // Offset type used for global offsets + using offset_t = detail::choose_signed_offset_t; + + // The lengths output value type + using length_t = cub::detail::non_void_value_t; + + // Generator type for providing 1s values for run-length reduction + using lengths_input_iterator_t = ::cuda::constant_iterator; + + using accum_t = ::cuda::std::__accumulator_t; + using key_t = cub::detail::non_void_value_t>; + using policy_selector_t = detail::rle::encode::policy_selector_from_types; + + return detail::reduce_by_key::dispatch_streaming( + d_temp_storage, + temp_storage_bytes, + d_in, + d_unique_out, + lengths_input_iterator_t(length_t{1}), + d_counts_out, + d_num_runs_out, + equality_op{}, + reduction_op{}, + static_cast(num_items), + stream, + __policy_selector_adapter{}); + } + + //! @rst + //! Computes a run-length encoding of the sequence ``d_in``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - For the *i*\ :sup:`th` run encountered, the first key of the run and + //! its length are written to ``d_unique_out[i]`` and ``d_counts_out[i]``, respectively. + //! - The total number of runs encountered is written to ``d_num_runs_out``. + //! - The ``==`` equality operator is used to determine whether values are equivalent + //! - In-place operations are not supported. There must be no overlap between any of the provided ranges: + //! + //! - ``[d_unique_out, d_unique_out + *d_num_runs_out)`` + //! - ``[d_counts_out, d_counts_out + *d_num_runs_out)`` + //! - ``[d_num_runs_out, d_num_runs_out + 1)`` + //! - ``[d_in, d_in + num_items)`` + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the run-length encoding of a sequence of ``int`` values + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_run_length_encode_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin encode-env + //! :end-before: example-end encode-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam UniqueOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing unique output items @iterator + //! + //! @tparam LengthsOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing output counts @iterator + //! + //! @tparam NumRunsOutputIteratorT + //! **[inferred]** Output iterator type for recording the number of runs encountered @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_in + //! Pointer to the input sequence of keys + //! + //! @param[out] d_unique_out + //! Pointer to the output sequence of unique keys (one key per run) + //! + //! @param[out] d_counts_out + //! Pointer to the output sequence of run-lengths (one count per run) + //! + //! @param[out] d_num_runs_out + //! Pointer to total number of runs + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Encode( + InputIteratorT d_in, + UniqueOutputIteratorT d_unique_out, + LengthsOutputIteratorT d_counts_out, + NumRunsOutputIteratorT d_num_runs_out, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceRunLengthEncode::Encode"); + + using equality_op = ::cuda::std::equal_to<>; + using reduction_op = ::cuda::std::plus<>; + using offset_t = detail::choose_signed_offset_t; + using length_t = cub::detail::non_void_value_t; + using lengths_input_iterator_t = ::cuda::constant_iterator; + using accum_t = ::cuda::std::__accumulator_t; + using key_t = cub::detail::non_void_value_t>; + using default_policy_selector = detail::rle::encode::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&]([[maybe_unused]] auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::reduce_by_key::dispatch_streaming( + storage, + bytes, + d_in, + d_unique_out, + lengths_input_iterator_t(length_t{1}), + d_counts_out, + d_num_runs_out, + equality_op{}, + reduction_op{}, + static_cast(num_items), + stream, + __policy_selector_adapter{}); + }); + } + + //! @rst + //! Enumerates the starting offsets and lengths of all non-trivial runs + //! (of ``length > 1``) of same-valued keys in the sequence ``d_in``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - For the *i*\ :sup:`th` non-trivial run, the run's starting offset and + //! its length are written to ``d_offsets_out[i]`` and ``d_lengths_out[i]``, respectively. + //! - The total number of runs encountered is written to ``d_num_runs_out``. + //! - The ``==`` equality operator is used to determine whether values are equivalent + //! - In-place operations are not supported. There must be no overlap between any of the provided ranges: + //! + //! - ``[d_offsets_out, d_offsets_out + *d_num_runs_out)`` + //! - ``[d_lengths_out, d_lengths_out + *d_num_runs_out)`` + //! - ``[d_num_runs_out, d_num_runs_out + 1)`` + //! - ``[d_in, d_in + num_items)`` + //! + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the identification of non-trivial runs + //! within a sequence of ``int`` values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [0, 2, 2, 9, 5, 5, 5, 8] + //! int *d_offsets_out; // e.g., [ , , , , , , , ] + //! int *d_lengths_out; // e.g., [ , , , , , , , ] + //! int *d_num_runs_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceRunLengthEncode::NonTrivialRuns( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_offsets_out, d_lengths_out, d_num_runs_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run encoding + //! cub::DeviceRunLengthEncode::NonTrivialRuns( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_offsets_out, d_lengths_out, d_num_runs_out, num_items); + //! + //! // d_offsets_out <-- [1, 4] + //! // d_lengths_out <-- [2, 3] + //! // d_num_runs_out <-- [2] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OffsetsOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing run-offset values @iterator + //! + //! @tparam LengthsOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing run-length values @iterator + //! + //! @tparam NumRunsOutputIteratorT + //! **[inferred]** Output iterator type for recording the number of runs encountered @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @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_in + //! Pointer to input sequence of data items + //! + //! @param[out] d_offsets_out + //! Pointer to output sequence of run-offsets + //! (one offset per non-trivial run) + //! + //! @param[out] d_lengths_out + //! Pointer to output sequence of run-lengths (one count per non-trivial run) + //! + //! @param[out] d_num_runs_out + //! Pointer to total number of runs (i.e., length of `d_offsets_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t NonTrivialRuns( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OffsetsOutputIteratorT d_offsets_out, + LengthsOutputIteratorT d_lengths_out, + NumRunsOutputIteratorT d_num_runs_out, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceRunLengthEncode::NonTrivialRuns"); + + using global_offset_t = detail::choose_signed_offset_t; + using equality_op = ::cuda::std::equal_to<>; + return detail::rle::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_offsets_out, + d_lengths_out, + d_num_runs_out, + equality_op{}, + static_cast(num_items), + stream); + } + + //! @rst + //! Enumerates the starting offsets and lengths of all non-trivial runs + //! (of ``length > 1``) of same-valued keys in the sequence ``d_in``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - For the *i*\ :sup:`th` non-trivial run, the run's starting offset and + //! its length are written to ``d_offsets_out[i]`` and ``d_lengths_out[i]``, respectively. + //! - The total number of runs encountered is written to ``d_num_runs_out``. + //! - The ``==`` equality operator is used to determine whether values are equivalent + //! - In-place operations are not supported. There must be no overlap between any of the provided ranges: + //! + //! - ``[d_offsets_out, d_offsets_out + *d_num_runs_out)`` + //! - ``[d_lengths_out, d_lengths_out + *d_num_runs_out)`` + //! - ``[d_num_runs_out, d_num_runs_out + 1)`` + //! - ``[d_in, d_in + num_items)`` + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the identification of non-trivial runs + //! within a sequence of ``int`` values using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_run_length_encode_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin non-trivial-runs-env + //! :end-before: example-end non-trivial-runs-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OffsetsOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing run-offset values @iterator + //! + //! @tparam LengthsOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing run-length values @iterator + //! + //! @tparam NumRunsOutputIteratorT + //! **[inferred]** Output iterator type for recording the number of runs encountered @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_in + //! Pointer to input sequence of data items + //! + //! @param[out] d_offsets_out + //! Pointer to output sequence of run-offsets + //! (one offset per non-trivial run) + //! + //! @param[out] d_lengths_out + //! Pointer to output sequence of run-lengths (one count per non-trivial run) + //! + //! @param[out] d_num_runs_out + //! Pointer to total number of runs (i.e., length of `d_offsets_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t NonTrivialRuns( + InputIteratorT d_in, + OffsetsOutputIteratorT d_offsets_out, + LengthsOutputIteratorT d_lengths_out, + NumRunsOutputIteratorT d_num_runs_out, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceRunLengthEncode::NonTrivialRuns"); + + using global_offset_t = detail::choose_signed_offset_t; + using equality_op = ::cuda::std::equal_to<>; + using length_t = detail::non_void_value_t; + using key_t = detail::it_value_t; + using default_policy_selector = detail::rle::non_trivial_runs::policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::rle::dispatch( + storage, + bytes, + d_in, + d_offsets_out, + d_lengths_out, + d_num_runs_out, + equality_op{}, + static_cast(num_items), + stream, + policy_selector); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_scan.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_scan.cuh new file mode 100644 index 00000000..e2716571 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_scan.cuh @@ -0,0 +1,3591 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceScan provides device-wide, parallel operations for computing a prefix scan across a sequence of data +//! items residing within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceScan provides device-wide, parallel operations for computing a +//! prefix scan across a sequence of data items residing within +//! device-accessible memory. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! Given a sequence of input elements and a binary reduction operator, a +//! `prefix scan `_ produces an output +//! sequence where each element is computed to be the reduction of the elements +//! occurring earlier in the input sequence. *Prefix sum* connotes a prefix scan +//! with the addition operator. The term *inclusive* indicates that the +//! *i*\ :sup:`th` output reduction incorporates the *i*\ :sup:`th` input. +//! The term *exclusive* indicates the *i*\ :sup:`th` input is not +//! incorporated into the *i*\ :sup:`th` output reduction. When the input and +//! output sequences are the same, the scan is performed in-place. +//! +//! In order to provide an efficient parallel implementation, the binary reduction operator must be associative. That +//! is, ``op(op(a, b), c)`` must be equivalent to ``op(a, op(b, c))`` for any input values ``a``, ``b``, and ``c``. +//! +//! As of CUB 1.0.1 (2013), CUB's device-wide scan APIs have implemented our +//! *"decoupled look-back"* algorithm for performing global prefix scan with +//! only a single pass through the input data, as described in our 2016 technical +//! report [1]_. The central idea is to leverage a small, constant factor of +//! redundant work in order to overlap the latencies of global prefix +//! propagation with local computation. As such, our algorithm requires only +//! ``~2*n*`` data movement (``n`` inputs are read, ``n`` outputs are written), and +//! typically proceeds at "memcpy" speeds. Our algorithm supports inplace operations. +//! +//! .. [1] Duane Merrill and Michael Garland. `Single-pass Parallel Prefix Scan with Decoupled Look-back +//! `_, +//! *NVIDIA Technical Report NVR-2016-002*, 2016. +//! +//! Usage Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @cdp_class{DeviceScan} +//! @determinism{not_guaranteed} +//! +//! Determinism +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! ``cub::DeviceScan`` supports the :ref:`determinism guarantees ` as follows; the +//! default is ``not_guaranteed``. +//! +//! - ``run_to_run`` is supported for integral types with a known CUDA binary operator, and for floating-point +//! types with ``cuda::std::plus``. The floating-point ``plus`` case engages a stable, fixed reduction order so +//! results are reproducible across runs on the same GPU. +//! - ``gpu_to_gpu`` is supported for integral types with a known CUDA binary operator. (These are exactly +//! associative, so the result is already identical across GPUs.) +//! - Other combinations under ``run_to_run``/``gpu_to_gpu`` are rejected at compile time. +//! +//! Performance +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @linear_performance{prefix scan} +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All non-ByKey algorithms in DeviceScan that accept an environment can be tuned by passing a custom :ref:`policy +//! selector ` that returns a :cpp:struct:`cub::ScanPolicy`, as shown in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin exclusive-sum-policy-selector +//! :end-before: example-end exclusive-sum-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin exclusive-sum-tuning +//! :end-before: example-end exclusive-sum-tuning +//! +//! All ByKey algorithms in DeviceScan that accept an environment can be tuned by passing a custom :ref:`policy +//! selector ` that returns a :cpp:struct:`cub::ScanByKeyPolicy`, as shown in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_by_key_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin exclusive-sum-by-key-policy-selector +//! :end-before: example-end exclusive-sum-by-key-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_by_key_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin exclusive-sum-by-key-tuning +//! :end-before: example-end exclusive-sum-by-key-tuning +//! +//! @endrst +struct DeviceScan +{ + //! @cond + template + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t scan_impl_determinism( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init, + NumItemsT num_items, + cudaStream_t stream, + PolicySelectorT policy_selector) + { + // Unsigned integer type for global offsets + using offset_t = detail::choose_offset_t; + return detail::scan::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + init, + static_cast(num_items), + stream, + policy_selector); + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t scan_impl_env( + InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op, InitValueT init, NumItemsT num_items, const EnvT& env) + { + static_assert(!::cuda::std::execution::__queryable_with, + "Determinism should be used inside requires to have an effect."); + + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + using accum_t = + ::cuda::std::__accumulator_t, + ::cuda::std::_If<::cuda::std::is_same_v, + cub::detail::it_value_t, + typename InitValueT::value_type>>; + using offset_t = detail::choose_offset_t; + + constexpr bool is_run_to_run_required = + ::cuda::std::is_same_v; + constexpr bool is_gpu_to_gpu_required = + ::cuda::std::is_same_v; + constexpr bool is_safe_integral_op = + ::cuda::std::is_integral_v && detail::is_cuda_binary_operator; + constexpr bool is_fp_plus_op = + ::cuda::std::is_floating_point_v && detail::is_cuda_std_plus_v; + + // run_to_run determinism is supported only with integral types with known operators, or floating-point types with + // plus operator + static_assert(!is_run_to_run_required || is_safe_integral_op || is_fp_plus_op, + "run_to_run deterministic scan requires either integral types with known operators, " + "or floating-point types with plus operator"); + + // gpu_to_gpu determinism is only supported with integral types with known operators + static_assert(!is_gpu_to_gpu_required || is_safe_integral_op, + "gpu_to_gpu deterministic scan requires integral types with known operators"); + + static constexpr bool stable_reduction_order = is_run_to_run_required && is_fp_plus_op; + + using default_policy_selector_t = detail::scan:: + policy_selector_from_types; + + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return scan_impl_determinism( + storage, bytes, d_in, d_out, scan_op, init, num_items, stream, policy_selector); + }); + } + + template + CUB_RUNTIME_FUNCTION static cudaError_t scan_by_key_impl( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + EqualityOpT equality_op, + ScanOpT scan_op, + InitValueT init_value, + NumItemsT num_items, + cudaStream_t stream) + { + using offset_t = detail::choose_offset_t; + using accum_t = ::cuda::std::__accumulator_t< + ScanOpT, + cub::detail::it_value_t, + ::cuda::std:: + _If<::cuda::std::is_same_v, cub::detail::it_value_t, InitValueT>>; + + using default_policy_selector_t = + detail::scan_by_key::policy_selector_from_types, + accum_t, + cub::detail::it_value_t, + ScanOpT>; + + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + + // we would not need to override the accumulator type, but we must ensure it's the same as for the policy here + return detail::scan_by_key::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + scan_op, + init_value, + static_cast(num_items), + stream, + policy_selector_t{}); + } + //! @endcond + + //! @name Exclusive scans + //! @{ + + //! @rst + //! Computes a device-wide exclusive prefix sum. + //! The value of ``0`` is applied as the initial value, and is assigned to ``*d_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. + //! The range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix sum of an ``int`` + //! device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_out; // e.g., [ , , , , , , ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::ExclusiveSum( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run exclusive prefix sum + //! cub::DeviceScan::ExclusiveSum( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, num_items); + //! + //! // d_out <-- [0, 8, 14, 21, 26, 29, 29] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveSum( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::ExclusiveSum"); + + // Unsigned integer type for global offsets + using OffsetT = detail::choose_offset_t; + using init_value_t = cub::detail::it_value_t; + + // Initial value + init_value_t init_value{}; + + return detail::scan::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + ::cuda::std::plus<>{}, + detail::InputValue(init_value), + static_cast(num_items), + stream); + } + + //! @rst + //! Computes a device-wide exclusive prefix sum. + //! The value of ``0`` is applied as the initial value, and is assigned to ``*d_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. + //! The range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a user-defined exclusive-scan of a + //! device vector of ``float`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-env-determinism + //! :end-before: example-end exclusive-sum-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + ExclusiveSum(InputIteratorT d_in, OutputIteratorT d_out, NumItemsT num_items, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::ExclusiveSum"); + + using init_value_t = cub::detail::it_value_t; + init_value_t init_value{}; + + return scan_impl_env( + d_in, d_out, ::cuda::std::plus<>{}, detail::InputValue(init_value), num_items, env); + } + + //! @rst + //! Computes a device-wide exclusive prefix sum in-place. + //! The value of ``0`` is applied as the initial value, and is assigned to ``*d_data``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix sum of an ``int`` + //! device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_data; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::ExclusiveSum( + //! d_temp_storage, temp_storage_bytes, + //! d_data, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run exclusive prefix sum + //! cub::DeviceScan::ExclusiveSum( + //! d_temp_storage, temp_storage_bytes, + //! d_data, num_items); + //! + //! // d_data <-- [0, 8, 14, 21, 26, 29, 29] + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading scan inputs and wrigin scan outputs + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_data`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveSum( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + return ExclusiveSum(d_temp_storage, temp_storage_bytes, d_data, d_data, num_items, stream); + } + + //! @rst + //! Computes a device-wide exclusive prefix sum in-place. + //! The value of ``0`` is applied as the initial value, and is assigned to ``*d_data``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates an in-place exclusive prefix sum of an + //! ``int`` device vector using a stream environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-inplace-env + //! :end-before: example-end exclusive-sum-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing scan data + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in,out] d_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_data`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + ExclusiveSum(IteratorT d_data, NumItemsT num_items, const EnvT& env = {}) + { + return ExclusiveSum(d_data, d_data, num_items, env); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan using the specified + //! binary associative ``scan_op`` functor. The ``init_value`` value is applied as + //! the initial value, and is assigned to ``*d_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix min-scan of an ``int`` device vector + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include // for INT_MAX + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_out; // e.g., [ , , , , , , ] + //! CustomMin min_op; + //! ... + //! + //! // Determine temporary device storage requirements for exclusive + //! // prefix scan + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::ExclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, min_op, (int) INT_MAX, num_items); + //! + //! // Allocate temporary storage for exclusive prefix scan + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run exclusive prefix min-scan + //! cub::DeviceScan::ExclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, min_op, (int) INT_MAX, num_items); + //! + //! // d_out <-- [2147483647, 8, 6, 6, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan (and is assigned to `*d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScan( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::ExclusiveScan"); + + // Unsigned integer type for global offsets + using OffsetT = detail::choose_offset_t; + + return detail::scan::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + detail::InputValue(init_value), + static_cast(num_items), + stream); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan using the specified + //! binary associative ``scan_op`` functor. The ``init_value`` value is applied as + //! the initial value, and is assigned to ``*d_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a user-defined exclusive-scan of a + //! device vector of ``float`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-env-determinism + //! :end-before: example-end exclusive-scan-env-determinism + //! + //! The code snippet below illustrates an exclusive-scan using a custom stream + //! and ``not_guaranteed`` determinism. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-env-stream + //! :end-before: example-end exclusive-scan-env-stream + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan (and is assigned to `*d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0, + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScan( + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::ExclusiveScan"); + + return scan_impl_env(d_in, d_out, scan_op, detail::InputValue(init_value), num_items, env); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan using the specified + //! binary associative ``scan_op`` functor. The ``init_value`` value is applied as + //! the initial value, and is assigned to ``*d_data``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix min-scan of an + //! ``int`` device vector: + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include // for INT_MAX + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_data; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! CustomMin min_op; + //! ... + //! + //! // Determine temporary device storage requirements for exclusive + //! // prefix scan + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::ExclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_data, min_op, (int) INT_MAX, num_items); + //! + //! // Allocate temporary storage for exclusive prefix scan + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run exclusive prefix min-scan + //! cub::DeviceScan::ExclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_data, min_op, (int) INT_MAX, num_items); + //! + //! // d_data <-- [2147483647, 8, 6, 6, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs and writing scan outputs + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan (and is assigned to `*d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScan( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + ScanOpT scan_op, + InitValueT init_value, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + return ExclusiveScan(d_temp_storage, temp_storage_bytes, d_data, d_data, scan_op, init_value, num_items, stream); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan in-place using the specified + //! binary associative ``scan_op`` functor. The ``init_value`` value is applied as + //! the initial value, and is assigned to ``*d_data``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates an in-place exclusive prefix scan of an + //! ``int`` device vector with a user-supplied initial value and a stream environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-inplace-env + //! :end-before: example-end exclusive-scan-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing scan data + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the ``init_value`` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in,out] d_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_data`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + ExclusiveScan(IteratorT d_data, ScanOpT scan_op, InitValueT init_value, NumItemsT num_items, const EnvT& env = {}) + { + return ExclusiveScan(d_data, d_data, scan_op, init_value, num_items, env); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan using the specified + //! binary associative ``scan_op`` functor. The ``init_value`` value is provided as a future value. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. + //! The range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix min-scan of an ``int`` device vector + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include // for INT_MAX + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_out; // e.g., [ , , , , , , ] + //! int *d_init_iter; // e.g., INT_MAX + //! CustomMin min_op; + //! + //! auto future_init_value = + //! cub::FutureValue(d_init_iter); + //! + //! ... + //! + //! // Determine temporary device storage requirements for exclusive + //! // prefix scan + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::ExclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, min_op, future_init_value, num_items); + //! + //! // Allocate temporary storage for exclusive prefix scan + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run exclusive prefix min-scan + //! cub::DeviceScan::ExclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, min_op, future_init_value, num_items); + //! + //! // d_out <-- [2147483647, 8, 6, 6, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan (and is assigned to `*d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScan( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + FutureValue init_value, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::ExclusiveScan"); + + // Unsigned integer type for global offsets + using OffsetT = detail::choose_offset_t; + + return detail::scan::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + detail::InputValue(init_value), + static_cast(num_items), + stream); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! The ``init_value`` value is provided as a future value. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix min-scan of an ``int`` device vector + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include // for INT_MAX + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_data; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_init_iter; // e.g., INT_MAX + //! CustomMin min_op; + //! + //! auto future_init_value = + //! cub::FutureValue(d_init_iter); + //! + //! ... + //! + //! // Determine temporary device storage requirements for exclusive + //! // prefix scan + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::ExclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_data, min_op, future_init_value, num_items); + //! + //! // Allocate temporary storage for exclusive prefix scan + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run exclusive prefix min-scan + //! cub::DeviceScan::ExclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_data, min_op, future_init_value, num_items); + //! + //! // d_data <-- [2147483647, 8, 6, 6, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs and writing scan outputs + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan (and is assigned to `*d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScan( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + ScanOpT scan_op, + FutureValue init_value, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + return ExclusiveScan(d_temp_storage, temp_storage_bytes, d_data, d_data, scan_op, init_value, num_items, stream); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan in-place using the specified + //! binary associative ``scan_op`` functor. The ``init_value`` value is provided as a future value. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates an in-place exclusive prefix scan of an + //! ``int`` device vector with a future initial value and a stream environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-future-inplace-env + //! :end-before: example-end exclusive-scan-future-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing scan data + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the ``init_value`` + //! + //! @tparam InitValueIterT + //! **[inferred]** Iterator type for the future value + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in,out] d_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan, provided as a future value + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_data`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScan( + IteratorT d_data, + ScanOpT scan_op, + FutureValue init_value, + NumItemsT num_items, + const EnvT& env = {}) + { + return ExclusiveScan(d_data, d_data, scan_op, init_value, num_items, env); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! The ``init_value`` value is provided as a future value. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. + //! The range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix min-scan of an ``int`` device vector + //! using a future value for the initial value. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-future-env + //! :end-before: example-end exclusive-scan-future-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam InitValueIterT + //! **[inferred]** Random-access iterator type used to access the initial value on device + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan (and is assigned to `*d_out`), + //! provided as a future value + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0, + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScan( + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + FutureValue init_value, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::ExclusiveScan"); + + return scan_impl_env( + d_in, d_out, scan_op, detail::InputValue(init_value), num_items, env); + } + + //! @} + + //! @name Inclusive scans + //! @{ + + //! @rst + //! Computes a device-wide inclusive prefix sum. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix sum of an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_out; // e.g., [ , , , , , , ] + //! ... + //! + //! // Determine temporary device storage requirements for inclusive + //! // prefix sum + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::InclusiveSum( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, num_items); + //! + //! // Allocate temporary storage for inclusive prefix sum + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run inclusive prefix sum + //! cub::DeviceScan::InclusiveSum( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, num_items); + //! + //! // d_out <-- [8, 14, 21, 26, 29, 29, 38] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t InclusiveSum( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::InclusiveSum"); + + // Unsigned integer type for global offsets + using OffsetT = detail::choose_offset_t; + + return detail::scan::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + ::cuda::std::plus<>{}, + NullType{}, + static_cast(num_items), + stream); + } + + //! @rst + //! Computes a device-wide inclusive prefix sum in-place. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix sum of an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_data; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! ... + //! + //! // Determine temporary device storage requirements for inclusive + //! // prefix sum + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::InclusiveSum( + //! d_temp_storage, temp_storage_bytes, + //! d_data, num_items); + //! + //! // Allocate temporary storage for inclusive prefix sum + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run inclusive prefix sum + //! cub::DeviceScan::InclusiveSum( + //! d_temp_storage, temp_storage_bytes, + //! d_data, num_items); + //! + //! // d_data <-- [8, 14, 21, 26, 29, 29, 38] + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs and writing scan outputs + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t InclusiveSum( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + return InclusiveSum(d_temp_storage, temp_storage_bytes, d_data, d_data, num_items, stream); + } + + //! @rst + //! Computes a device-wide inclusive prefix sum in-place. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates an in-place inclusive prefix sum of an + //! ``int`` device vector using a stream environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-sum-inplace-env + //! :end-before: example-end inclusive-sum-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing scan data + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in,out] d_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_data`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + InclusiveSum(IteratorT d_data, NumItemsT num_items, const EnvT& env = {}) + { + return InclusiveSum(d_data, d_data, num_items, env); + } + + //! @rst + //! Computes a device-wide inclusive prefix sum. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. + //! The range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix sum of an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-sum-env-determinism + //! :end-before: example-end inclusive-sum-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + InclusiveSum(InputIteratorT d_in, OutputIteratorT d_out, NumItemsT num_items, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::InclusiveSum"); + + return scan_impl_env(d_in, d_out, ::cuda::std::plus<>{}, NullType{}, num_items, env); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix min-scan of an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include // for INT_MAX + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_out; // e.g., [ , , , , , , ] + //! CustomMin min_op; + //! ... + //! + //! // Determine temporary device storage requirements for inclusive + //! // prefix scan + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::InclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, min_op, num_items); + //! + //! // Allocate temporary storage for inclusive prefix scan + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run inclusive prefix min-scan + //! cub::DeviceScan::InclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, min_op, num_items); + //! + //! // d_out <-- [8, 6, 6, 5, 3, 0, 0] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t InclusiveScan( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::InclusiveScan"); + + // Unsigned integer type for global offsets + using OffsetT = detail::choose_offset_t; + + return detail::scan::dispatch( + d_temp_storage, temp_storage_bytes, d_in, d_out, scan_op, NullType(), static_cast(num_items), stream); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! The result of applying the ``scan_op`` binary operator to ``init_value`` value and ``*d_in`` + //! is assigned to ``*d_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive max-scan of an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin device-inclusive-scan + //! :end-before: example-end device-inclusive-scan + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to the size in bytes of the `d_temp_storage` allocation + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the inclusive scan (`scan_op(init_value, d_in[0])` + //! is assigned to `*d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! CUDA stream to launch kernels within. + template + CUB_RUNTIME_FUNCTION static cudaError_t InclusiveScanInit( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::InclusiveScanInit"); + + // Unsigned integer type for global offsets + using OffsetT = detail::choose_offset_t; + + return detail::scan::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + detail::InputValue(init_value), + static_cast(num_items), + stream); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix min-scan of an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include // for INT_MAX + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_data; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! CustomMin min_op; + //! ... + //! + //! // Determine temporary device storage requirements for inclusive + //! // prefix scan + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::InclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_data, min_op, num_items); + //! + //! // Allocate temporary storage for inclusive prefix scan + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run inclusive prefix min-scan + //! cub::DeviceScan::InclusiveScan( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, min_op, num_items); + //! + //! // d_data <-- [8, 6, 6, 5, 3, 0, 0] + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs and writing scan outputs + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t InclusiveScan( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + ScanOpT scan_op, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + return InclusiveScan(d_temp_storage, temp_storage_bytes, d_data, d_data, scan_op, num_items, stream); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan in-place using the specified + //! binary associative ``scan_op`` functor. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates an in-place inclusive prefix scan of an + //! ``int`` device vector using a stream environment. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-inplace-env + //! :end-before: example-end inclusive-scan-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing scan data + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in,out] d_data + //! Random-access iterator to the sequence of data items + //! + //! @param[in] scan_op + //! Binary scan functor + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_data`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + InclusiveScan(IteratorT d_data, ScanOpT scan_op, NumItemsT num_items, const EnvT& env = {}) + { + return InclusiveScan(d_data, d_data, scan_op, num_items, env); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix sum of an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-env + //! :end-before: example-end inclusive-scan-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + InclusiveScan(InputIteratorT d_in, OutputIteratorT d_out, ScanOpT scan_op, NumItemsT num_items, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::InclusiveScan"); + + return scan_impl_env(d_in, d_out, scan_op, NullType{}, num_items, env); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! The result of applying the ``scan_op`` binary operator to ``init_value`` value and ``*d_in`` + //! is assigned to ``*d_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix sum of an ``int`` device vector + //! with an initial value. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-init-env + //! :end-before: example-end inclusive-scan-init-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the inclusive scan (`scan_op(init_value, d_in[0])` + //! is assigned to `*d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t InclusiveScanInit( + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::InclusiveScanInit"); + + return scan_impl_env( + d_in, d_out, scan_op, detail::InputValue(init_value), num_items, env); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! The result of applying the ``scan_op`` binary operator to ``init_value`` value and ``*d_in`` + //! is assigned to ``*d_out``. + //! + //! .. versionadded:: 3.5.0 + //! First appears in CUDA Toolkit 13.5. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix sum of an ``int`` device vector + //! with an initial value. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-future-init-env + //! :end-before: example-end inclusive-scan-future-init-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueIterT + //! **[inferred]** Random-access iterator type used to access the initial value on device + //! + //! @tparam InitValueBoundsT + //! **[inferred]** Static bounds on `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the inclusive scan (`scan_op(init_value, d_in[0])` + //! is assigned to `*d_out`), provided as deferred value. The deferred value must model an + //! iterator (i.e. the contained value should be dereferenceable to a value). + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t InclusiveScanInit( + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + const ::cuda::args::deferred& init_value, + NumItemsT num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::InclusiveScanInit"); + + static_assert(::cuda::std::indirectly_readable, "The deferred value must model an iterator"); + using __init_value_type = typename ::cuda::std::remove_cvref_t::__element_type; + + auto __fut = FutureValue<__init_value_type, InitValueIterT>{::cuda::args::__unwrap(init_value)}; + + return scan_impl_env( + d_in, d_out, scan_op, detail::InputValue<__init_value_type, InitValueIterT>(__fut), num_items, env); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan using the specified binary associative ``scan_op`` functor. + //! The result of applying the ``scan_op`` binary operator to ``init_value`` value and ``*d_in`` + //! is assigned to ``*d_out``. + //! + //! .. versionadded:: 3.5.0 + //! First appears in CUDA Toolkit 13.5. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - When ``d_in`` and ``d_out`` are equal, the scan is performed in-place. The + //! range ``[d_in, d_in + num_items)`` and ``[d_out, d_out + num_items)`` + //! shall not overlap in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive max-scan of an ``int`` device vector + //! with the initial value provided as a deferred device value. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin device-inclusive-scan-init-deferred + //! :end-before: example-end device-inclusive-scan-init-deferred + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan inputs @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueIterT + //! **[inferred]** Random-access iterator type used to access the initial value on device + //! + //! @tparam InitValueBoundsT + //! **[inferred]** Static bounds on `init_value` + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to the size in bytes of the `d_temp_storage` allocation + //! + //! @param[in] d_in + //! Random-access iterator to the input sequence of data items + //! + //! @param[out] d_out + //! Random-access iterator to the output sequence of data items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the inclusive scan (`scan_op(init_value, d_in[0])` + //! is assigned to `*d_out`), provided as deferred value. The deferred value must model an + //! iterator (i.e. the contained value should be dereferenceable to a value). + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_in`) + //! + //! @param[in] stream + //! CUDA stream to launch kernels within. + template + CUB_RUNTIME_FUNCTION static cudaError_t InclusiveScanInit( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + const ::cuda::args::deferred& init_value, + NumItemsT num_items, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::InclusiveScanInit"); + + static_assert(::cuda::std::indirectly_readable, "The deferred value must model an iterator"); + using __init_value_type = typename ::cuda::std::remove_cvref_t::__element_type; + + // Unsigned integer type for global offsets + using OffsetT = detail::choose_offset_t; + + auto __fut = FutureValue<__init_value_type, InitValueIterT>{::cuda::args::__unwrap(init_value)}; + + return detail::scan::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + detail::InputValue<__init_value_type, InitValueIterT>(__fut), + static_cast(num_items), + stream); + } + + //! @} + + //! @name Scans by key + //! @{ + + //! @rst + //! Computes a device-wide exclusive prefix sum-by-key with key equality + //! defined by ``equality_op``. The value of ``0`` is applied as the initial + //! value, and is assigned to the beginning of each segment in ``d_values_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - ``d_keys_in`` may equal ``d_values_out`` but the range + //! ``[d_keys_in, d_keys_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - ``d_values_in`` may equal ``d_values_out`` but the range + //! ``[d_values_in, d_values_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix sum-by-key of an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_keys_in; // e.g., [0, 0, 1, 1, 1, 2, 2] + //! int *d_values_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_values_out; // e.g., [ , , , , , , ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::ExclusiveSumByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, d_values_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run exclusive prefix sum + //! cub::DeviceScan::ExclusiveSumByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, d_values_out, num_items); + //! + //! // d_values_out <-- [0, 8, 0, 7, 12, 0, 0] + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan keys inputs @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan values inputs @iterator + //! + //! @tparam ValuesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan values outputs @iterator + //! + //! @tparam EqualityOpT + //! **[inferred]** Functor type having member + //! `T operator()(const T &a, const T &b)` for binary operations that defines the equality of keys + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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 input iterator to the input sequence of key items + //! + //! @param[in] d_values_in + //! Random-access input iterator to the input sequence of value items + //! + //! @param[out] d_values_out + //! Random-access output iterator to the output sequence of value items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_keys_in` and `d_values_in`) + //! + //! @param[in] equality_op + //! Binary functor that defines the equality of keys. + //! Default is cuda::std::equal_to<>{}. + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template , + typename NumItemsT = uint32_t> + CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveSumByKey( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + NumItemsT num_items, + EqualityOpT equality_op = EqualityOpT(), + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::ExclusiveSumByKey"); + using init_value_t = cub::detail::it_value_t; + init_value_t init_value{}; + return scan_by_key_impl<::cuda::std::execution::env<>>( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + ::cuda::std::plus<>{}, + init_value, + num_items, + stream); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan-by-key using the + //! specified binary associative ``scan_op`` functor. The key equality is defined by + //! ``equality_op``. The ``init_value`` value is applied as the initial + //! value, and is assigned to the beginning of each segment in ``d_values_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - ``d_keys_in`` may equal ``d_values_out`` but the range + //! ``[d_keys_in, d_keys_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - ``d_values_in`` may equal ``d_values_out`` but the range + //! ``[d_values_in, d_values_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix min-scan-by-key of an ``int`` device vector + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include // for INT_MAX + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // CustomEqual functor + //! struct CustomEqual + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return a == b; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_keys_in; // e.g., [0, 0, 1, 1, 1, 2, 2] + //! int *d_values_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_values_out; // e.g., [ , , , , , , ] + //! CustomMin min_op; + //! CustomEqual equality_op; + //! ... + //! + //! // Determine temporary device storage requirements for exclusive + //! // prefix scan + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::ExclusiveScanByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, d_values_out, min_op, + //! (int) INT_MAX, num_items, equality_op); + //! + //! // Allocate temporary storage for exclusive prefix scan + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run exclusive prefix min-scan + //! cub::DeviceScan::ExclusiveScanByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, d_values_out, min_op, + //! (int) INT_MAX, num_items, equality_op); + //! + //! // d_values_out <-- [2147483647, 8, 2147483647, 7, 5, 2147483647, 0] + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan keys inputs @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan values inputs @iterator + //! + //! @tparam ValuesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan values outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam EqualityOpT + //! **[inferred]** Functor type having member + //! `T operator()(const T &a, const T &b)` for binary operations that defines the equality of keys + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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 input iterator to the input sequence of key items + //! + //! @param[in] d_values_in + //! Random-access input iterator to the input sequence of value items + //! + //! @param[out] d_values_out + //! Random-access output iterator to the output sequence of value items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan (and is assigned to the + //! beginning of each segment in `d_values_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_keys_in` and + //! `d_values_in`) + //! + //! @param[in] equality_op + //! Binary functor that defines the equality of keys. + //! Default is cuda::std::equal_to<>{}. + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template , + typename NumItemsT = uint32_t> + CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScanByKey( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + ScanOpT scan_op, + InitValueT init_value, + NumItemsT num_items, + EqualityOpT equality_op = EqualityOpT(), + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::ExclusiveScanByKey"); + return scan_by_key_impl<::cuda::std::execution::env<>>( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + scan_op, + init_value, + num_items, + stream); + } + + //! @rst + //! Computes a device-wide inclusive prefix sum-by-key with key equality defined by ``equality_op``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - ``d_keys_in`` may equal ``d_values_out`` but the range + //! ``[d_keys_in, d_keys_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - ``d_values_in`` may equal ``d_values_out`` but the range + //! ``[d_values_in, d_values_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix sum-by-key of an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_keys_in; // e.g., [0, 0, 1, 1, 1, 2, 2] + //! int *d_values_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_values_out; // e.g., [ , , , , , , ] + //! ... + //! + //! // Determine temporary device storage requirements for inclusive prefix sum + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::InclusiveSumByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, d_values_out, num_items); + //! + //! // Allocate temporary storage for inclusive prefix sum + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run inclusive prefix sum + //! cub::DeviceScan::InclusiveSumByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, d_values_out, num_items); + //! + //! // d_out <-- [8, 14, 7, 12, 15, 0, 9] + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan keys inputs @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan values inputs @iterator + //! + //! @tparam ValuesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan values outputs @iterator + //! + //! @tparam EqualityOpT + //! **[inferred]** Functor type having member + //! `T operator()(const T &a, const T &b)` for binary operations that defines the equality of keys + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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 input iterator to the input sequence of key items + //! + //! @param[in] d_values_in + //! Random-access input iterator to the input sequence of value items + //! + //! @param[out] d_values_out + //! Random-access output iterator to the output sequence of value items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_keys_in` and `d_values_in`) + //! + //! @param[in] equality_op + //! Binary functor that defines the equality of keys. + //! Default is cuda::std::equal_to<>{}. + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template , + typename NumItemsT = uint32_t> + CUB_RUNTIME_FUNCTION static cudaError_t InclusiveSumByKey( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + NumItemsT num_items, + EqualityOpT equality_op = EqualityOpT(), + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::InclusiveSumByKey"); + return scan_by_key_impl<::cuda::std::execution::env<>>( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + ::cuda::std::plus<>{}, + NullType{}, + num_items, + stream); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan-by-key using the + //! specified binary associative ``scan_op`` functor. The key equality is defined by ``equality_op``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - ``d_keys_in`` may equal ``d_values_out`` but the range + //! ``[d_keys_in, d_keys_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - ``d_values_in`` may equal ``d_values_out`` but the range + //! ``[d_values_in, d_values_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix min-scan-by-key of an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! #include // for INT_MAX + //! + //! // CustomMin functor + //! struct CustomMin + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return (b < a) ? b : a; + //! } + //! }; + //! + //! // CustomEqual functor + //! struct CustomEqual + //! { + //! template + //! __host__ __device__ __forceinline__ + //! T operator()(const T &a, const T &b) const { + //! return a == b; + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // input and output + //! int num_items; // e.g., 7 + //! int *d_keys_in; // e.g., [0, 0, 1, 1, 1, 2, 2] + //! int *d_values_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_values_out; // e.g., [ , , , , , , ] + //! CustomMin min_op; + //! CustomEqual equality_op; + //! ... + //! + //! // Determine temporary device storage requirements for inclusive prefix scan + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceScan::InclusiveScanByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, d_values_out, min_op, num_items, equality_op); + //! + //! // Allocate temporary storage for inclusive prefix scan + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run inclusive prefix min-scan + //! cub::DeviceScan::InclusiveScanByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, d_values_out, min_op, num_items, equality_op); + //! + //! // d_out <-- [8, 6, 7, 5, 3, 0, 0] + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan keys inputs @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan values inputs @iterator + //! + //! @tparam ValuesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan values outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam EqualityOpT + //! **[inferred]** Functor type having member + //! `T operator()(const T &a, const T &b)` for binary operations that defines the equality of keys + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @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 input iterator to the input sequence of key items + //! + //! @param[in] d_values_in + //! Random-access input iterator to the input sequence of value items + //! + //! @param[out] d_values_out + //! Random-access output iterator to the output sequence of value items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_keys_in` and `d_values_in`) + //! + //! @param[in] equality_op + //! Binary functor that defines the equality of keys. + //! Default is cuda::std::equal_to<>{}. + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template , + typename NumItemsT = uint32_t> + CUB_RUNTIME_FUNCTION static cudaError_t InclusiveScanByKey( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + ScanOpT scan_op, + NumItemsT num_items, + EqualityOpT equality_op = EqualityOpT(), + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceScan::InclusiveScanByKey"); + return scan_by_key_impl<::cuda::std::execution::env<>>( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + scan_op, + NullType{}, + num_items, + stream); + } + + //! @rst + //! Computes a device-wide exclusive prefix sum-by-key with key equality + //! defined by ``equality_op``. The value of ``0`` is applied as the initial + //! value, and is assigned to the beginning of each segment in ``d_values_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - ``d_keys_in`` may equal ``d_values_out`` but the range + //! ``[d_keys_in, d_keys_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - ``d_values_in`` may equal ``d_values_out`` but the range + //! ``[d_values_in, d_values_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix sum-by-key of an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_by_key_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-sum-by-key-env + //! :end-before: example-end exclusive-sum-by-key-env + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan keys inputs @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan values inputs @iterator + //! + //! @tparam ValuesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan values outputs @iterator + //! + //! @tparam EqualityOpT + //! **[inferred]** Functor type having member + //! `T operator()(const T &a, const T &b)` for binary operations that defines the equality of keys + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_keys_in + //! Random-access input iterator to the input sequence of key items + //! + //! @param[in] d_values_in + //! Random-access input iterator to the input sequence of value items + //! + //! @param[out] d_values_out + //! Random-access output iterator to the output sequence of value items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_keys_in` and `d_values_in`) + //! + //! @param[in] equality_op + //! Binary functor that defines the equality of keys. + //! Default is cuda::std::equal_to<>{}. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t< + !::cuda::std::is_same_v && !::cuda::std::is_null_pointer_v + && !::cuda::std::is_same_v, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveSumByKey( + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + NumItemsT num_items, + EqualityOpT equality_op = EqualityOpT(), + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::ExclusiveSumByKey"); + + using init_value_t = cub::detail::it_value_t; + return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + using tuning_t = decltype(tuning); + return scan_by_key_impl( + storage, + bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + ::cuda::std::plus<>{}, + init_value_t{}, + num_items, + stream); + }); + } + + //! @rst + //! Computes a device-wide exclusive prefix scan-by-key using the + //! specified binary associative ``scan_op`` functor. The key equality is defined by + //! ``equality_op``. The ``init_value`` value is applied as the initial + //! value, and is assigned to the beginning of each segment in ``d_values_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - ``d_keys_in`` may equal ``d_values_out`` but the range + //! ``[d_keys_in, d_keys_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - ``d_values_in`` may equal ``d_values_out`` but the range + //! ``[d_values_in, d_values_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the exclusive prefix scan-by-key of an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_by_key_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin exclusive-scan-by-key-env + //! :end-before: example-end exclusive-scan-by-key-env + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan keys inputs @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan values inputs @iterator + //! + //! @tparam ValuesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan values outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam InitValueT + //! **[inferred]** Type of the `init_value` + //! + //! @tparam EqualityOpT + //! **[inferred]** Functor type having member + //! `T operator()(const T &a, const T &b)` for binary operations that defines the equality of keys + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_keys_in + //! Random-access input iterator to the input sequence of key items + //! + //! @param[in] d_values_in + //! Random-access input iterator to the input sequence of value items + //! + //! @param[out] d_values_out + //! Random-access output iterator to the output sequence of value items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] init_value + //! Initial value to seed the exclusive scan (and is assigned to the + //! beginning of each segment in `d_values_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_keys_in` and `d_values_in`) + //! + //! @param[in] equality_op + //! Binary functor that defines the equality of keys. + //! Default is cuda::std::equal_to<>{}. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t< + !::cuda::std::is_same_v && !::cuda::std::is_null_pointer_v + && !::cuda::std::is_same_v, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ExclusiveScanByKey( + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + ScanOpT scan_op, + InitValueT init_value, + NumItemsT num_items, + EqualityOpT equality_op = EqualityOpT(), + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::ExclusiveScanByKey"); + + return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + using tuning_t = decltype(tuning); + return scan_by_key_impl( + storage, bytes, d_keys_in, d_values_in, d_values_out, equality_op, scan_op, init_value, num_items, stream); + }); + } + + //! @rst + //! Computes a device-wide inclusive prefix sum-by-key with key equality defined by ``equality_op``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative sum operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - ``d_keys_in`` may equal ``d_values_out`` but the range + //! ``[d_keys_in, d_keys_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - ``d_values_in`` may equal ``d_values_out`` but the range + //! ``[d_values_in, d_values_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix sum-by-key of an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_by_key_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-sum-by-key-env + //! :end-before: example-end inclusive-sum-by-key-env + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan keys inputs @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan values inputs @iterator + //! + //! @tparam ValuesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan values outputs @iterator + //! + //! @tparam EqualityOpT + //! **[inferred]** Functor type having member + //! `T operator()(const T &a, const T &b)` for binary operations that defines the equality of keys + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_keys_in + //! Random-access input iterator to the input sequence of key items + //! + //! @param[in] d_values_in + //! Random-access input iterator to the input sequence of value items + //! + //! @param[out] d_values_out + //! Random-access output iterator to the output sequence of value items + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_keys_in` and `d_values_in`) + //! + //! @param[in] equality_op + //! Binary functor that defines the equality of keys. + //! Default is cuda::std::equal_to<>{}. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t< + !::cuda::std::is_same_v && !::cuda::std::is_null_pointer_v + && !::cuda::std::is_same_v, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t InclusiveSumByKey( + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + NumItemsT num_items, + EqualityOpT equality_op = EqualityOpT(), + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::InclusiveSumByKey"); + + return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + using tuning_t = decltype(tuning); + return scan_by_key_impl( + storage, + bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + ::cuda::std::plus<>{}, + NullType{}, + num_items, + stream); + }); + } + + //! @rst + //! Computes a device-wide inclusive prefix scan-by-key using the + //! specified binary associative ``scan_op`` functor. The key equality is defined by ``equality_op``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Supports non-commutative scan operators. + //! - Results are not deterministic for pseudo-associative operators (e.g., + //! addition of floating-point types). Results for pseudo-associative + //! operators may vary from run to run. Additional details can be found in + //! the @lookback description. + //! - ``d_keys_in`` may equal ``d_values_out`` but the range + //! ``[d_keys_in, d_keys_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! - ``d_values_in`` may equal ``d_values_out`` but the range + //! ``[d_values_in, d_values_in + num_items)`` and the range + //! ``[d_values_out, d_values_out + num_items)`` shall not overlap otherwise. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the inclusive prefix scan-by-key of an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_scan_by_key_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin inclusive-scan-by-key-env + //! :end-before: example-end inclusive-scan-by-key-env + //! + //! @endrst + //! + //! @tparam KeysInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan keys inputs @iterator + //! + //! @tparam ValuesInputIteratorT + //! **[inferred]** Random-access input iterator type for reading scan values inputs @iterator + //! + //! @tparam ValuesOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing scan values outputs @iterator + //! + //! @tparam ScanOpT + //! **[inferred]** Binary associative scan functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam EqualityOpT + //! **[inferred]** Functor type having member + //! `T operator()(const T &a, const T &b)` for binary operations that defines the equality of keys + //! + //! @tparam NumItemsT + //! **[inferred]** An integral type representing the number of input elements + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type providing stream, memory resource, + //! or determinism requirements. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_keys_in + //! Random-access input iterator to the input sequence of key items + //! + //! @param[in] d_values_in + //! Random-access input iterator to the input sequence of value items + //! + //! @param[out] d_values_out + //! Random-access output iterator to the output sequence of value items + //! + //! @param[in] scan_op + //! Binary associative scan functor + //! + //! @param[in] num_items + //! Total number of input items (i.e., the length of `d_keys_in` and `d_values_in`) + //! + //! @param[in] equality_op + //! Binary functor that defines the equality of keys. + //! Default is cuda::std::equal_to<>{}. + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template , + typename NumItemsT = uint32_t, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t< + !::cuda::std::is_same_v && !::cuda::std::is_null_pointer_v + && !::cuda::std::is_same_v, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t InclusiveScanByKey( + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + ScanOpT scan_op, + NumItemsT num_items, + EqualityOpT equality_op = EqualityOpT(), + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceScan::InclusiveScanByKey"); + + return detail::dispatch_with_env(env, [&]([[maybe_unused]] auto tuning, void* storage, size_t& bytes, auto stream) { + using tuning_t = decltype(tuning); + return scan_by_key_impl( + storage, bytes, d_keys_in, d_values_in, d_values_out, equality_op, scan_op, NullType{}, num_items, stream); + }); + } + + //! @} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_radix_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_radix_sort.cuh new file mode 100644 index 00000000..9ff1da24 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_radix_sort.cuh @@ -0,0 +1,2621 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceSegmentedRadixSort provides device-wide, parallel operations for computing a batched radix sort across +//! multiple, non-overlapping sequences of data items residing within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceSegmentedRadixSort provides device-wide, parallel operations +//! for computing a batched radix sort across multiple, non-overlapping +//! sequences of data items residing within device-accessible memory. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The `radix sorting method `_ +//! arranges items into ascending (or descending) order. The algorithm relies +//! upon a positional representation for keys, i.e., each key is comprised of an +//! ordered sequence of symbols (e.g., digits, characters, etc.) specified from +//! least-significant to most-significant. For a given input sequence of keys +//! and a set of rules specifying a total ordering of the symbolic alphabet, the +//! radix sorting method produces a lexicographic ordering of those keys. +//! +//! See Also +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! DeviceSegmentedRadixSort shares its implementation with DeviceRadixSort. See +//! that algorithm's documentation for more information. +//! +//! Segments are not required to be contiguous. Any element of input(s) or +//! output(s) outside the specified segments will not be accessed nor modified. +//! +//! Usage Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @cdp_class{DeviceSegmentedRadixSort} +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceSegmentedRadixSort that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::SegmentedRadixSortPolicy`, as shown +//! in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin segmented-radix-sort-policy-selector +//! :end-before: example-end segmented-radix-sort-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin segmented-radix-sort-keys-tuning +//! :end-before: example-end segmented-radix-sort-keys-tuning +//! +//! @endrst +struct DeviceSegmentedRadixSort +{ +private: + // Name reported for NVTX ranges + _CCCL_HOST_DEVICE static constexpr auto GetName() -> const char* + { + return "cub::DeviceSegmentedRadixSort"; + } + +public: + //! @name Key-value pairs + //! @{ + + //! @rst + //! Sorts segments of key-value pairs into ascending order. (``~2N`` auxiliary storage required) + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - @devicestorageNP For sorting using only ``O(P)`` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``int`` keys with associated vector of ``int`` values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_values_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedRadixSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedRadixSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [6, 7, 8, 0, 3, 5, 9] + //! // d_values_out <-- [1, 2, 0, 5, 4, 3, 6] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! Random-access input iterator to the sequence of beginning offsets of + //! length `num_segments`, such that `d_begin_offsets[i]` is the first + //! element of the *i*th data segment in `d_keys_*` and `d_values_*` + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. If + //! ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + + return detail::segmented_radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + false, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into ascending order. (``~N`` auxiliary storage required) + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits specified and the targeted device architecture). + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter is + //! specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and yield + //! a corresponding performance improvement. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of `int` keys with associated vector of ``int`` values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! int *d_value_buf; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_value_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a set of DoubleBuffers to wrap pairs of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! cub::DoubleBuffer d_values(d_value_buf, d_value_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedRadixSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedRadixSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [6, 7, 8, 0, 3, 5, 9] + //! // d_values.Current() <-- [5, 4, 3, 1, 2, 0, 6] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + return detail::segmented_radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + true, + stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Sorts segments of key-value pairs into ascending order. (``~2N`` auxiliary storage required) + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-radix-sort-pairs-env + //! :end-before: example-end segmented-radix-sort-pairs-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! Random-access input iterator to the sequence of beginning offsets of + //! length `num_segments`, such that `d_begin_offsets[i]` is the first + //! element of the *i*th data segment in `d_keys_*` and `d_values_*` + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. If + //! ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using SegmentSizeT = ::cuda::std::int32_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::segmented_radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + static_cast<::cuda::std::int64_t>(num_items), + static_cast<::cuda::std::int64_t>(num_segments), + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + false, + stream, + /* decomposer */ {}, + tuning_env); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Sorts segments of key-value pairs into ascending order. (``~N`` auxiliary storage required) + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits specified and the targeted device architecture). + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter is + //! specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-radix-sort-pairs-db-env + //! :end-before: example-end segmented-radix-sort-pairs-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using SegmentSizeT = ::cuda::std::int32_t; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::segmented_radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + true, + stream, + /* decomposer */ {}, + tuning_env); + }); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. (``~2N`` auxiliary storage required). + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter is + //! specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and `out` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - @devicestorageNP For sorting using only ``O(P)`` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``int`` keys with associated vector of ``int`` values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_values_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedRadixSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedRadixSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [8, 7, 6, 9, 5, 3, 0] + //! // d_values_out <-- [0, 2, 1, 6, 3, 4, 5] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + + return detail::segmented_radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + false, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. (``~N`` auxiliary storage required). + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits specified and the targeted device architecture). + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter is + //! specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``int`` keys with associated vector of ``int`` values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! int *d_value_buf; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_value_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a set of DoubleBuffers to wrap pairs of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! cub::DoubleBuffer d_values(d_value_buf, d_value_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedRadixSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedRadixSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [8, 7, 6, 9, 5, 3, 0] + //! // d_values.Current() <-- [0, 2, 1, 6, 3, 4, 5] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + return detail::segmented_radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + true, + stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Sorts segments of key-value pairs into descending order. (``~2N`` auxiliary storage required) + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter is + //! specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-radix-sort-pairs-descending-env + //! :end-before: example-end segmented-radix-sort-pairs-descending-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using SegmentSizeT = ::cuda::std::int32_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::segmented_radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + false, + stream, + /* decomposer */ {}, + tuning_env); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Sorts segments of key-value pairs into descending order. (``~N`` auxiliary storage required) + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits specified and the targeted device architecture). + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter is + //! specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-radix-sort-pairs-descending-db-env + //! :end-before: example-end segmented-radix-sort-pairs-descending-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer + //! contains the unsorted input values and, upon return, is updated to point + //! to the sorted output values + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using SegmentSizeT = ::cuda::std::int32_t; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::segmented_radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + true, + stream, + /* decomposer */ {}, + tuning_env); + }); + } + + //! @} + //! @name Keys-only + //! @{ + + //! @rst + //! Sorts segments of keys into ascending order. (``~2N`` auxiliary storage required) + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter + //! is specified as ``segment_offsets + 1``). + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - @devicestorageNP For sorting using only ``O(P)`` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of `int` keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedRadixSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedRadixSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [6, 7, 8, 0, 3, 5, 9] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + // Null value type + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values; + + return detail::segmented_radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + false, + stream); + } + + //! @rst + //! Sorts segments of keys into ascending order. (``~N`` auxiliary storage required). + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter + //! is specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``int`` keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedRadixSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedRadixSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [6, 7, 8, 0, 3, 5, 9] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) + //! needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + // Null value type + DoubleBuffer d_values; + + return detail::segmented_radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + true, + stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Sorts segments of keys into ascending order. (``~2N`` auxiliary storage required) + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter + //! is specified as ``segment_offsets + 1``). + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-radix-sort-keys-env + //! :end-before: example-end segmented-radix-sort-keys-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using SegmentSizeT = ::cuda::std::int32_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::segmented_radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + false, + stream, + /* decomposer */ {}, + tuning_env); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Sorts segments of keys into ascending order. (``~N`` auxiliary storage required) + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-radix-sort-keys-db-env + //! :end-before: example-end segmented-radix-sort-keys-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items within the segmented array + //! + //! @param[in] num_segments + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + // Null value type + DoubleBuffer d_values; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::segmented_radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + true, + stream, + /* decomposer */ {}, + tuning_env); + }); + } + + //! @rst + //! Sorts segments of keys into descending order. (``~2N`` auxiliary storage required). + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter + //! is specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - @devicestorageNP For sorting using only ``O(P)`` temporary storage, see + //! the sorting interface using DoubleBuffer wrappers below. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``int`` keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedRadixSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedRadixSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [8, 7, 6, 9, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., sizeof(unsigned int) * 8) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values; + + return detail::segmented_radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + false, + stream); + } + + //! @rst + //! Sorts segments of keys into descending order. (``~N`` auxiliary storage required). + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! - @devicestorageP + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of `int` keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedRadixSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedRadixSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [8, 7, 6, 9, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key + //! comparison (e.g., `sizeof(unsigned int) * 8`) + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + // Null value type + DoubleBuffer d_values; + + return detail::segmented_radix_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + true, + stream); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Sorts segments of keys into descending order. (``~2N`` auxiliary storage required) + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter + //! is specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-radix-sort-keys-descending-env + //! :end-before: example-end segmented-radix-sort-keys-descending-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_keys_in + //! Pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + using SegmentSizeT = ::cuda::std::int32_t; + + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::segmented_radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + false, + stream, + /* decomposer */ {}, + tuning_env); + }); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! Sorts segments of keys into descending order. (``~N`` auxiliary storage required) + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! DoubleBuffer structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the + //! number of key bits specified and the targeted device architecture). + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - An optional bit subrange ``[begin_bit, end_bit)`` of differentiating key + //! bits can be specified. This can reduce overall sorting overhead and + //! yield a corresponding performance improvement. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! - Note, the size of any segment may not exceed ``INT_MAX``. Please consider using ``DeviceSegmentedSort`` instead, + //! if the size of at least one of your segments could exceed ``INT_MAX``. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_radix_sort_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-radix-sort-keys-descending-db-env + //! :end-before: example-end segmented-radix-sort-keys-descending-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items within the segmented array, including items not + //! covered by segments. `num_items` should match the largest element within + //! the range `[d_end_offsets, d_end_offsets + num_segments)`. + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] begin_bit + //! **[optional]** The least-significant bit index (inclusive) needed for key comparison + //! + //! @param[in] end_bit + //! **[optional]** The most-significant bit index (exclusive) needed for key comparison + //! (e.g., ``sizeof(unsigned int) * 8``) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit = 0, + int end_bit = sizeof(KeyT) * 8, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + + // Signed integer type for global offsets + using SegmentSizeT = ::cuda::std::int32_t; + + // Null value type + DoubleBuffer d_values; + + return detail::dispatch_with_env(env, [&](auto tuning_env, void* storage, size_t& bytes, auto stream) { + return detail::segmented_radix_sort::dispatch( + storage, + bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + true, + stream, + /* decomposer */ {}, + tuning_env); + }); + } + + //! @} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_reduce.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_reduce.cuh new file mode 100644 index 00000000..2b86f4a0 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_reduce.cuh @@ -0,0 +1,2736 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceSegmentedReduce provides device-wide, parallel operations for computing a batched reduction across +//! multiple sequences of data items residing within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceSegmentedReduce provides device-wide, parallel operations for +//! computing a reduction across multiple sequences of data items +//! residing within device-accessible memory. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! A `reduction `_ +//! (or *fold*) uses a binary combining operator to compute a single aggregate +//! from a sequence of input elements. +//! +//! Usage Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @cdp_class{DeviceSegmentedReduce} +//! @determinism{run_to_run} +//! +//! Determinism +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! ``cub::DeviceSegmentedReduce`` supports ``not_guaranteed`` and ``run_to_run`` (default ``run_to_run``). +//! ``gpu_to_gpu`` is not supported and is rejected at compile time. See the +//! :ref:`determinism guarantees ` for what each level means. +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceSegmentedReduce that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::SegmentedReducePolicy`, as shown in +//! the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin segmented-reduce-sum-policy-selector +//! :end-before: example-end segmented-reduce-sum-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin segmented-reduce-sum-tuning +//! :end-before: example-end segmented-reduce-sum-tuning +//! +//! @endrst +struct DeviceSegmentedReduce +{ +private: + template , + typename InputIteratorT, + typename OutputIteratorT> + CUB_RUNTIME_FUNCTION static cudaError_t fixed_size_arg_impl( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + cudaStream_t stream) + { + // `offset_t` a.k.a `SegmentSizeT` is fixed to `int` type now, but later can be changed to accept + // integral constant or larger integral types + using offset_t = int; + + using input_value_t = cub::detail::it_value_t; + using output_tuple_t = cub::detail::non_void_value_t>; + using accum_t = output_tuple_t; + using init_value_t = detail::reduce::empty_problem_init_t; + using output_key_t = typename output_tuple_t::first_type; + using output_value_t = typename output_tuple_t::second_type; + + static_assert(::cuda::std::is_same_v, "Output key type must be int."); + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + auto d_indexed_in = THRUST_NS_QUALIFIER::make_transform_iterator( + THRUST_NS_QUALIFIER::counting_iterator<::cuda::std::int64_t>{0}, + detail::segmented_reduce::generate_idx_value(d_in, segment_size)); + using arg_index_input_iterator_t = decltype(d_indexed_in); + + constexpr bool is_min = ::cuda::std::is_same_v; + auto sentinel = + is_min ? ::cuda::std::numeric_limits::max() : ::cuda::std::numeric_limits::lowest(); + init_value_t initial_value{accum_t(1, sentinel)}; + + using default_policy_selector_t = + detail::segmented_reduce::policy_selector_from_types; + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + + return detail::segmented_reduce::dispatch_fixed_size( + d_temp_storage, + temp_storage_bytes, + d_indexed_in, + d_out, + num_segments, + static_cast(segment_size), + ReductionOpT(), + initial_value, + stream, + policy_selector_t{}); + } + + template + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t fixed_size_arg_impl_env( + InputIteratorT d_in, OutputIteratorT d_out, ::cuda::std::int64_t num_segments, int segment_size, const EnvT& env) + { + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported for device segmented reductions "); + + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return fixed_size_arg_impl( + d_temp_storage, temp_storage_bytes, d_in, d_out, num_segments, segment_size, stream); + }); + } + + template + CUB_RUNTIME_FUNCTION static cudaError_t variable_size_env_impl( + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + ReductionOpT reduction_op, + InitValueT initial_value, + const EnvT& env) + { + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported for device segmented reductions "); + + static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); + if constexpr (::cuda::std::is_integral_v) + { + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + using default_policy_selector_t = + detail::segmented_reduce::policy_selector_from_types; + using policy_selector_t = ::cuda::std::execution:: + __query_result_or_t; + // TODO: in most cases we can just take the default AccumT and OffsetT. Refactor this + return detail::segmented_reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + reduction_op, + initial_value, + 0, // max_segment_size + stream, + policy_selector_t{}); + }); + } + _CCCL_UNREACHABLE(); + } + + template + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t fixed_size_impl( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + ReductionOpT reduction_op, + T initial_value, + cudaStream_t stream = nullptr) + { + // `offset_t` a.k.a `SegmentSizeT` is fixed to `int` type now, but later can be changed to accept + // integral constant or larger integral types + using offset_t = int; + + return detail::segmented_reduce::dispatch_fixed_size( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + static_cast(segment_size), + reduction_op, + initial_value, + stream); + } + + template + CUB_RUNTIME_FUNCTION static cudaError_t fixed_size_env_impl( + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + OffsetT segment_size, + ReductionOpT reduction_op, + InitValueT initial_value, + const EnvT& env) + { + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported for device segmented reductions "); + + using default_policy_selector_t = + detail::segmented_reduce::policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return detail::segmented_reduce::dispatch_fixed_size( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + segment_size, + reduction_op, + initial_value, + stream, + policy_selector); + }); + } + +public: + //! @rst + //! Computes a device-wide segmented reduction using the specified + //! binary ``reduction_op`` functor. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Does not support binary reduction operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a custom min-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-custommin + //! :end-before: example-end segmented-reduce-custommin + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-reduce + //! :end-before: example-end segmented-reduce-reduce + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam T + //! **[inferred]** Data element type that is convertible to the `value` type of `InputIteratorT` + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] initial_value + //! Initial value of the reduction for each segment + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t Reduce( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + ReductionOpT reduction_op, + T initial_value, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::Reduce"); + + using OffsetT = detail::common_iterator_value_t; + static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); + if constexpr (::cuda::std::is_integral_v) + { + return detail::segmented_reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + reduction_op, + initial_value, // zero-initialize + 0, // max_segment_size + stream); + } + _CCCL_UNREACHABLE(); + } + + //! @rst + //! Computes a device-wide segmented reduction using the specified + //! binary ``reduction_op`` functor. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Does not support binary reduction operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a custom min-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-reduce-env + //! :end-before: example-end segmented-reduce-reduce-env + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-reduce-env-determinism + //! :end-before: example-end segmented-reduce-reduce-env-determinism + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam T + //! **[inferred]** Data element type that is convertible to the `value` type of `InputIteratorT` + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] initial_value + //! Initial value of the reduction for each segment + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t Reduce( + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + ReductionOpT reduction_op, + T initial_value, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Reduce"); + + using OffsetT = detail::common_iterator_value_t; + using AccumT = ::cuda::std::__accumulator_t, T>; + + return variable_size_env_impl( + d_in, d_out, num_segments, d_begin_offsets, d_end_offsets, reduction_op, initial_value, env); + } + + //! @rst + //! Computes a device-wide segmented reduction using the specified + //! binary ``reduction_op`` functor and a fixed segment size. + //! + //! .. versionadded:: 3.2.0 + //! First appears in CUDA Toolkit 13.2. + //! + //! - Does not support binary reduction operators that are non-commutative. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates a custom min-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-custommin + //! :end-before: example-end segmented-reduce-custommin + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-reduce + //! :end-before: example-end fixed-size-segmented-reduce-reduce + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam T + //! **[inferred]** Data element type that is convertible to the `value` type of `InputIteratorT` + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregates + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] initial_value + //! Initial value of the reduction for each segment + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t Reduce( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + ReductionOpT reduction_op, + T initial_value, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::Reduce"); + return fixed_size_impl( + d_temp_storage, temp_storage_bytes, d_in, d_out, num_segments, segment_size, reduction_op, initial_value, stream); + } + + //! @rst + //! Computes a device-wide segmented reduction using the specified + //! binary ``reduction_op`` functor and a fixed segment size. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Does not support binary reduction operators that are non-commutative. + //! - Provides "run-to-run" determinism for pseudo-associative reduction + //! (e.g., addition of floating point types) on the same GPU device. + //! However, results for pseudo-associative reduction may be inconsistent + //! from one device to a another device of a different compute-capability + //! because CUB can employ different tile-sizing for different architectures. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-reduce-env + //! :end-before: example-end fixed-size-segmented-reduce-reduce-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam ReductionOpT + //! **[inferred]** Binary reduction functor type having member `T operator()(const T &a, const T &b)` + //! + //! @tparam T + //! **[inferred]** Data element type that is convertible to the `value` type of `InputIteratorT` + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregates + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] reduction_op + //! Binary reduction functor + //! + //! @param[in] initial_value + //! Initial value of the reduction for each segment + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t Reduce( + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + ReductionOpT reduction_op, + T initial_value, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Reduce"); + + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported for device segmented reductions "); + + // `offset_t` a.k.a `SegmentSizeT` is fixed to `int` type now, but later can be changed to accept + // integral constant or larger integral types + using offset_t = int; + using accum_t = ::cuda::std::__accumulator_t, T>; + + return fixed_size_env_impl( + d_in, d_out, num_segments, static_cast(segment_size), reduction_op, initial_value, env); + } + + //! @rst + //! Computes a device-wide segmented sum using the addition (``+``) operator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``0`` as the initial value of the reduction for each segment. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Does not support ``+`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the sum reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-sum + //! :end-before: example-end segmented-reduce-sum + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments`, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template ::value_type, + typename ::cuda::std::iterator_traits::value_type>> + CUB_RUNTIME_FUNCTION static cudaError_t + Sum(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::Sum"); + + using OffsetT = detail::common_iterator_value_t; + using OutputT = detail::non_void_value_t>; + using init_value_t = OutputT; + static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); + if constexpr (::cuda::std::is_integral_v) + { + return detail::segmented_reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + ::cuda::std::plus<>{}, + init_value_t{}, // zero-initialize + 0, // max_segment_size + stream); + } + _CCCL_UNREACHABLE(); + } + + //! @rst + //! Computes a device-wide segmented sum using the addition (``+``) operator. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Uses ``0`` as the initial value of the reduction for each segment. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Does not support ``+`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the sum reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-sum-env + //! :end-before: example-end segmented-reduce-sum-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments`, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Sum(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Sum"); + + using OffsetT = detail::common_iterator_value_t; + using OutputT = detail::non_void_value_t>; + using init_value_t = OutputT; + using op_t = ::cuda::std::plus<>; + using AccumT = ::cuda::std::__accumulator_t, init_value_t>; + + return variable_size_env_impl( + d_in, d_out, num_segments, d_begin_offsets, d_end_offsets, op_t{}, init_value_t{}, env); + } + + //! @rst + //! Computes a device-wide segmented sum using the addition (``+``) operator. + //! + //! .. versionadded:: 3.2.0 + //! First appears in CUDA Toolkit 13.2. + //! + //! - Uses ``0`` as the initial value of the reduction for each segment. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the sum reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-sum + //! :end-before: example-end fixed-size-segmented-reduce-sum + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t + Sum(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + cudaStream_t stream = nullptr) + { + static_assert(!::cuda::std::is_same_v, + "InputIteratorT must be a real iterator; void* has no iterator_traits::value_type."); + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::Sum"); + using init_value_t = detail::non_void_value_t>; + return fixed_size_impl( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + segment_size, + ::cuda::std::plus{}, + init_value_t{}, + stream); + } + + //! @rst + //! Computes a device-wide segmented sum using the addition (``+``) operator + //! and a fixed segment size. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Uses ``0`` as the initial value of the reduction for each segment. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-sum-env + //! :end-before: example-end fixed-size-segmented-reduce-sum-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Sum(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Sum"); + + using output_t = detail::non_void_value_t>; + + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported for device segmented reductions "); + + // `offset_t` a.k.a `SegmentSizeT` is fixed to `int` type now, but later can be changed to accept + // integral constant or larger integral types + using offset_t = int; + using op_t = ::cuda::std::plus<>; + using accum_t = ::cuda::std::__accumulator_t, output_t>; + + return fixed_size_env_impl( + d_in, d_out, num_segments, static_cast(segment_size), op_t{}, output_t{}, env); + } + + //! @rst + //! Computes a device-wide segmented minimum using the less-than (``<``) operator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``::cuda::std::numeric_limits::max()`` as the initial value of the reduction for each segment. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter is + //! specified as ``segment_offsets + 1``). + //! - Does not support ``<`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the min-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-min + //! :end-before: example-end segmented-reduce-min + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template ::value_type, + typename ::cuda::std::iterator_traits::value_type>> + CUB_RUNTIME_FUNCTION static cudaError_t + Min(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::Min"); + + using OffsetT = detail::common_iterator_value_t; + using InputT = detail::it_value_t; + using init_value_t = InputT; + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); + if constexpr (::cuda::std::is_integral_v) + { + return detail::segmented_reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + ::cuda::minimum<>{}, + ::cuda::std::numeric_limits::max(), + 0, // max_segment_size + stream); + } + _CCCL_UNREACHABLE(); + } + + //! @rst + //! Computes a device-wide segmented minimum using the less-than (``<``) operator. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Uses ``::cuda::std::numeric_limits::max()`` as the initial value of the reduction for each segment. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter is + //! specified as ``segment_offsets + 1``). + //! - Does not support ``<`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the min-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-min-env + //! :end-before: example-end segmented-reduce-min-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Min(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Min"); + + using OffsetT = detail::common_iterator_value_t; + using InputT = detail::it_value_t; + using init_value_t = InputT; + using op_t = ::cuda::minimum<>; + using AccumT = ::cuda::std::__accumulator_t, init_value_t>; + + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + return variable_size_env_impl( + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + op_t{}, + ::cuda::std::numeric_limits::max(), + env); + } + + //! @rst + //! Computes a device-wide segmented minimum using the less-than (``<``) operator. + //! + //! .. versionadded:: 3.2.0 + //! First appears in CUDA Toolkit 13.2. + //! + //! - Uses ``::cuda::std::numeric_limits::max()`` as the initial value of the reduction for each segment. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the min-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-min + //! :end-before: example-end fixed-size-segmented-reduce-min + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t + Min(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::Min"); + using input_t = detail::it_value_t; + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + return fixed_size_impl( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + segment_size, + ::cuda::minimum<>{}, + ::cuda::std::numeric_limits::max(), + stream); + } + + //! @rst + //! Computes a device-wide segmented minimum using the less-than (``<``) operator + //! and a fixed segment size. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Uses ``::cuda::std::numeric_limits::max()`` as the initial value of the reduction for each segment. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-min-env + //! :end-before: example-end fixed-size-segmented-reduce-min-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Min(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Min"); + + using input_t = detail::it_value_t; + + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported for device segmented reductions "); + + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + // `offset_t` a.k.a `SegmentSizeT` is fixed to `int` type now, but later can be changed to accept + // integral constant or larger integral types + using offset_t = int; + using op_t = ::cuda::minimum<>; + using accum_t = ::cuda::std::__accumulator_t, input_t>; + + return fixed_size_env_impl( + d_in, + d_out, + num_segments, + static_cast(segment_size), + op_t{}, + ::cuda::std::numeric_limits::max(), + env); + } + + //! @rst + //! Finds the first device-wide minimum in each segment using the + //! less-than (``<``) operator, also returning the in-segment index of that item. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The output value type of ``d_out`` is ``cub::KeyValuePair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The minimum of the *i*\ :sup:`th` segment is written to + //! ``d_out[i].value`` and its offset in that segment is written to ``d_out[i].key``. + //! - The ``{1, ::cuda::std::numeric_limits::max()}`` tuple is produced for zero-length inputs + //! + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased for both + //! the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where the latter + //! is specified as ``segment_offsets + 1``). + //! - Does not support ``<`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmin-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-argmin + //! :end-before: example-end segmented-reduce-argmin + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `KeyValuePair`) @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template ::value_type, + typename ::cuda::std::iterator_traits::value_type>> + CUB_RUNTIME_FUNCTION static cudaError_t ArgMin( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::ArgMin"); + + // Using common iterator value type is a breaking change, see: + // https://github.com/NVIDIA/cccl/pull/414#discussion_r1330632615 + using OverrideOffsetT = int; // detail::common_iterator_value_t; + + using InputValueT = detail::it_value_t; + using OutputTupleT = detail::non_void_value_t>; + using OutputKeyT = typename OutputTupleT::Key; + using OutputValueT = typename OutputTupleT::Value; + using OverrideAccumT = OutputTupleT; + using init_value_t = detail::reduce::empty_problem_init_t; + + static_assert(::cuda::std::is_same_v, "Output key type must be int."); + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + // Wrapped input iterator to produce index-value tuples + using ArgIndexInputIteratorT = ArgIndexInputIterator; + ArgIndexInputIteratorT d_indexed_in(d_in); + + init_value_t initial_value{OverrideAccumT(1, ::cuda::std::numeric_limits::max())}; + + static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); + if constexpr (::cuda::std::is_integral_v) + { + return detail::segmented_reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_indexed_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + cub::ArgMin{}, + initial_value, + 0, // max_segment_size + stream); + } + _CCCL_UNREACHABLE(); + } + + //! @rst + //! Finds the first device-wide minimum in each segment using the + //! less-than (``<``) operator, also returning the in-segment index of that item. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The output value type of ``d_out`` is ``cub::KeyValuePair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The minimum of the *i*\ :sup:`th` segment is written to + //! ``d_out[i].value`` and its offset in that segment is written to ``d_out[i].key``. + //! - The ``{1, ::cuda::std::numeric_limits::max()}`` tuple is produced for zero-length inputs + //! + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Does not support ``<`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmin-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-argmin-env + //! :end-before: example-end segmented-reduce-argmin-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `cub::KeyValuePair`) @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ArgMin( + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::ArgMin"); + + // Using common iterator value type is a breaking change, see: + // https://github.com/NVIDIA/cccl/pull/414#discussion_r1330632615 + using OverrideOffsetT = int; // detail::common_iterator_value_t; + + using InputValueT = detail::it_value_t; + using OutputTupleT = detail::non_void_value_t>; + using OutputKeyT = typename OutputTupleT::Key; + using OutputValueT = typename OutputTupleT::Value; + using OverrideAccumT = OutputTupleT; + using init_value_t = detail::reduce::empty_problem_init_t; + + static_assert(::cuda::std::is_same_v, "Output key type must be int."); + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + // Wrapped input iterator to produce index-value tuples + using ArgIndexInputIteratorT = ArgIndexInputIterator; + ArgIndexInputIteratorT d_indexed_in(d_in); + + init_value_t initial_value{OverrideAccumT(1, ::cuda::std::numeric_limits::max())}; + + return variable_size_env_impl( + d_indexed_in, d_out, num_segments, d_begin_offsets, d_end_offsets, cub::ArgMin{}, initial_value, env); + } + + //! @rst + //! Finds the first device-wide minimum in each segment using the + //! less-than (``<``) operator, also returning the in-segment index of that item. + //! + //! .. versionadded:: 3.2.0 + //! First appears in CUDA Toolkit 13.2. + //! + //! - The output value type of ``d_out`` is ``::cuda::std::pair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The minimum of the *i*\ :sup:`th` segment is written to + //! ``d_out[i].second`` and its offset in that segment is written to ``d_out[i].first``. + //! - The ``{1, ::cuda::std::numeric_limits::max()}`` tuple is produced for zero-length inputs + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmin-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-argmin + //! :end-before: example-end fixed-size-segmented-reduce-argmin + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `cuda::std::pair`) @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t ArgMin( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::ArgMin"); + return fixed_size_arg_impl( + d_temp_storage, temp_storage_bytes, d_in, d_out, num_segments, segment_size, stream); + } + + //! @rst + //! Finds the first device-wide minimum in each segment using the + //! less-than (``<``) operator, also returning the in-segment index of that item, + //! with a fixed segment size. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The output value type of ``d_out`` is ``::cuda::std::pair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The minimum of the *i*\ :sup:`th` segment is written to + //! ``d_out[i].second`` and its offset in that segment is written to ``d_out[i].first``. + //! - The ``{1, ::cuda::std::numeric_limits::max()}`` tuple is produced for zero-length inputs + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-argmin-env + //! :end-before: example-end fixed-size-segmented-reduce-argmin-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `cuda::std::pair`) @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + ArgMin(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::ArgMin"); + return fixed_size_arg_impl_env(d_in, d_out, num_segments, segment_size, env); + } + + //! @rst + //! Computes a device-wide segmented maximum using the greater-than (``>``) operator. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Uses ``::cuda::std::numeric_limits::lowest()`` as the initial value of the reduction. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Does not support ``>`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the max-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-max + //! :end-before: example-end segmented-reduce-max + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template ::value_type, + typename ::cuda::std::iterator_traits::value_type>> + CUB_RUNTIME_FUNCTION static cudaError_t + Max(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::Max"); + + using OffsetT = detail::common_iterator_value_t; + using InputT = cub::detail::it_value_t; + using init_value_t = InputT; + + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); + if constexpr (::cuda::std::is_integral_v) + { + return detail::segmented_reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + ::cuda::maximum<>{}, + ::cuda::std::numeric_limits::lowest(), + 0, // max_segment_size + stream); + } + _CCCL_UNREACHABLE(); + } + + //! @rst + //! Computes a device-wide segmented maximum using the greater-than (``>``) operator. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Uses ``::cuda::std::numeric_limits::lowest()`` as the initial value of the reduction. + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Does not support ``>`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the max-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-max-env + //! :end-before: example-end segmented-reduce-max-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Max(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Max"); + + using OffsetT = detail::common_iterator_value_t; + using InputT = cub::detail::it_value_t; + using init_value_t = InputT; + using op_t = ::cuda::maximum<>; + using AccumT = ::cuda::std::__accumulator_t, init_value_t>; + + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + return variable_size_env_impl( + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + op_t{}, + ::cuda::std::numeric_limits::lowest(), + env); + } + + //! @rst + //! Computes a device-wide segmented maximum using the greater-than (``>``) operator. + //! + //! .. versionadded:: 3.2.0 + //! First appears in CUDA Toolkit 13.2. + //! + //! - Uses ``::cuda::std::numeric_limits::lowest()`` as the initial value of the reduction. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the max-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-max + //! :end-before: example-end fixed-size-segmented-reduce-max + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t + Max(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::Max"); + using input_t = detail::it_value_t; + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + return fixed_size_impl( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + segment_size, + ::cuda::maximum<>{}, + ::cuda::std::numeric_limits::lowest(), + stream); + } + + //! @rst + //! Computes a device-wide segmented maximum using the greater-than (``>``) operator + //! and a fixed segment size. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - Uses ``::cuda::std::numeric_limits::lowest()`` as the initial value of the reduction. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-max-env + //! :end-before: example-end fixed-size-segmented-reduce-max-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + Max(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::Max"); + + using input_t = detail::it_value_t; + + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + + static_assert(!::cuda::std::is_same_v, + "gpu_to_gpu determinism is not supported for device segmented reductions "); + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + // `offset_t` a.k.a `SegmentSizeT` is fixed to `int` type now, but later can be changed to accept + // integral constant or larger integral types + using offset_t = int; + using op_t = ::cuda::maximum<>; + using accum_t = ::cuda::std::__accumulator_t, input_t>; + return fixed_size_env_impl( + d_in, + d_out, + num_segments, + static_cast(segment_size), + op_t{}, + ::cuda::std::numeric_limits::lowest(), + env); + } + + //! @rst + //! Finds the first device-wide maximum in each segment using the + //! greater-than (``>``) operator, also returning the in-segment index of that item + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The output value type of ``d_out`` is ``cub::KeyValuePair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The maximum of the *i*\ :sup:`th` segment is written to + //! ``d_out[i].value`` and its offset in that segment is written to ``d_out[i].key``. + //! - The ``{1, ::cuda::std::numeric_limits::lowest()}`` tuple is produced for zero-length inputs + //! + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Does not support ``>`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmax-reduction of a device vector + //! of `int` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-argmax + //! :end-before: example-end segmented-reduce-argmax + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items + //! (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `KeyValuePair`) @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length `num_segments`, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template ::value_type, + typename ::cuda::std::iterator_traits::value_type>> + CUB_RUNTIME_FUNCTION static cudaError_t ArgMax( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::ArgMax"); + + // Using common iterator value type is a breaking change, see: + // https://github.com/NVIDIA/cccl/pull/414#discussion_r1330632615 + using OverrideOffsetT = int; // detail::common_iterator_value_t; + + using InputValueT = cub::detail::it_value_t; + using OutputTupleT = cub::detail::non_void_value_t>; + using OverrideAccumT = OutputTupleT; + using init_value_t = detail::reduce::empty_problem_init_t; + using OutputKeyT = typename OutputTupleT::Key; + using OutputValueT = typename OutputTupleT::Value; + + static_assert(::cuda::std::is_same_v, "Output key type must be int."); + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + // Wrapped input iterator to produce index-value tuples + using ArgIndexInputIteratorT = ArgIndexInputIterator; + ArgIndexInputIteratorT d_indexed_in(d_in); + + init_value_t initial_value{OverrideAccumT(1, ::cuda::std::numeric_limits::lowest())}; + + static_assert(::cuda::std::is_integral_v, "Offset iterator value type should be integral."); + if constexpr (::cuda::std::is_integral_v) + { + return detail::segmented_reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_indexed_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + cub::ArgMax{}, + initial_value, + 0, // max_segment_size + stream); + } + _CCCL_UNREACHABLE(); + } + + //! @rst + //! Finds the first device-wide maximum in each segment using the + //! greater-than (``>``) operator, also returning the in-segment index of that item + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The output value type of ``d_out`` is ``cub::KeyValuePair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The maximum of the *i*\ :sup:`th` segment is written to + //! ``d_out[i].value`` and its offset in that segment is written to ``d_out[i].key``. + //! - The ``{1, ::cuda::std::numeric_limits::lowest()}`` tuple is produced for zero-length inputs + //! + //! - When input a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - Does not support ``>`` operators that are non-commutative. + //! - Let ``s`` be in ``[0, num_segments)``. The range + //! ``[d_out + d_begin_offsets[s], d_out + d_end_offsets[s])`` shall not + //! overlap ``[d_in + d_begin_offsets[s], d_in + d_end_offsets[s])``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)``. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmax-reduction of a device vector of ``int`` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-reduce-argmax-env + //! :end-before: example-end segmented-reduce-argmax-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `cub::KeyValuePair`) @iterator + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_in`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_in``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the *i*\ :sup:`th` is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t ArgMax( + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::ArgMax"); + + // Using common iterator value type is a breaking change, see: + // https://github.com/NVIDIA/cccl/pull/414#discussion_r1330632615 + using OverrideOffsetT = int; // detail::common_iterator_value_t; + + using InputValueT = cub::detail::it_value_t; + using OutputTupleT = cub::detail::non_void_value_t>; + using OverrideAccumT = OutputTupleT; + using init_value_t = detail::reduce::empty_problem_init_t; + using OutputKeyT = typename OutputTupleT::Key; + using OutputValueT = typename OutputTupleT::Value; + + static_assert(::cuda::std::is_same_v, "Output key type must be int."); + static_assert(::cuda::std::numeric_limits::is_specialized, + "numeric_limits must be specialized for the input value type"); + + // Wrapped input iterator to produce index-value tuples + using ArgIndexInputIteratorT = ArgIndexInputIterator; + ArgIndexInputIteratorT d_indexed_in(d_in); + + init_value_t initial_value{OverrideAccumT(1, ::cuda::std::numeric_limits::lowest())}; + + return variable_size_env_impl( + d_indexed_in, d_out, num_segments, d_begin_offsets, d_end_offsets, cub::ArgMax{}, initial_value, env); + } + + //! @rst + //! Finds the first device-wide maximum in each segment using the + //! greater-than (``>``) operator, also returning the in-segment index of that item + //! + //! .. versionadded:: 3.2.0 + //! First appears in CUDA Toolkit 13.2. + //! + //! - The output value type of ``d_out`` is ``::cuda::std::pair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The maximum of the *i*\ :sup:`th` segment is written to + //! ``d_out[i].second`` and its offset in that segment is written to ``d_out[i].first``. + //! - The ``{1, ::cuda::std::numeric_limits::lowest()}`` tuple is produced for zero-length inputs + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the argmax-reduction of a device vector + //! of `int` data elements. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-argmax + //! :end-before: example-end fixed-size-segmented-reduce-argmax + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items + //! (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `cuda::std::pair`) @iterator + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t ArgMax( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSegmentedReduce::ArgMax"); + return fixed_size_arg_impl( + d_temp_storage, temp_storage_bytes, d_in, d_out, num_segments, segment_size, stream); + } + + //! @rst + //! Finds the first device-wide maximum in each segment using the + //! greater-than (``>``) operator, also returning the in-segment index of that item, + //! with a fixed segment size. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The output value type of ``d_out`` is ``::cuda::std::pair`` + //! (assuming the value type of ``d_in`` is ``T``) + //! + //! - The maximum of the *i*\ :sup:`th` segment is written to + //! ``d_out[i].second`` and its offset in that segment is written to ``d_out[i].first``. + //! - The ``{1, ::cuda::std::numeric_limits::lowest()}`` tuple is produced for zero-length inputs + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_reduce_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin fixed-size-segmented-reduce-argmax-env + //! :end-before: example-end fixed-size-segmented-reduce-argmax-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items (of some type `T`) @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Output iterator type for recording the reduced aggregate + //! (having value type `cuda::std::pair`) @iterator + //! + //! @tparam EnvT + //! **[inferred]** Execution environment type. Default is ``cuda::std::execution::env<>``. + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output aggregate + //! + //! @param[in] num_segments + //! The number of segments that comprise the segmented reduction data + //! + //! @param[in] segment_size + //! The fixed segment size of each segment + //! + //! @param[in] env + //! @rst + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t + ArgMax(InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + int segment_size, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSegmentedReduce::ArgMax"); + return fixed_size_arg_impl_env(d_in, d_out, num_segments, segment_size, env); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_sort.cuh new file mode 100644 index 00000000..2974a2ec --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_segmented_sort.cuh @@ -0,0 +1,4681 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceSegmentedSort provides device-wide, parallel operations for computing a batched sort across multiple, +//! non-overlapping sequences of data items residing within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceSegmentedSort provides device-wide, parallel operations for +//! computing a batched sort across multiple, non-overlapping sequences of +//! data items residing within device-accessible memory. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The algorithm arranges items into ascending (or descending) order. +//! The underlying sorting algorithm is undefined. Depending on the segment size, +//! it might be radix sort, merge sort or something else. Therefore, no +//! assumptions on the underlying implementation should be made. +//! +//! Differences from DeviceSegmentedRadixSort +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! DeviceSegmentedRadixSort is optimized for significantly large segments (tens +//! of thousands of items and more). Nevertheless, some domains produce a wide +//! range of segment sizes. DeviceSegmentedSort partitions segments into size +//! groups and specialize sorting algorithms for each group. This approach leads +//! to better resource utilization in the presence of segment size imbalance or +//! moderate segment sizes (up to thousands of items). +//! This algorithm is more complex and consists of multiple kernels. This fact +//! leads to longer compilation times as well as larger binaries sizes. +//! +//! Supported Types +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! The algorithm has to satisfy the underlying algorithms restrictions. Radix +//! sort usage restricts the list of supported types. Therefore, +//! DeviceSegmentedSort can sort 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. +//! +//! Segments are not required to be contiguous. Any element of input(s) or +//! output(s) outside the specified segments will not be accessed nor modified. +//! +//! A simple example +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! .. code-block:: c++ +//! +//! #include +//! // or equivalently +//! +//! // Declare, allocate, and initialize device-accessible pointers +//! // for sorting data +//! int num_items; // e.g., 7 +//! int num_segments; // e.g., 3 +//! int *d_offsets; // e.g., [0, 3, 3, 7] +//! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] +//! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] +//! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] +//! int *d_values_out; // e.g., [-, -, -, -, -, -, -] +//! ... +//! +//! // Determine temporary device storage requirements +//! void *d_temp_storage = nullptr; +//! size_t temp_storage_bytes = 0; +//! cub::DeviceSegmentedSort::SortPairs( +//! d_temp_storage, temp_storage_bytes, +//! d_keys_in, d_keys_out, d_values_in, d_values_out, +//! num_items, num_segments, d_offsets, d_offsets + 1); +//! +//! // Allocate temporary storage +//! cudaMalloc(&d_temp_storage, temp_storage_bytes); +//! +//! // Run sorting operation +//! cub::DeviceSegmentedSort::SortPairs( +//! d_temp_storage, temp_storage_bytes, +//! d_keys_in, d_keys_out, d_values_in, d_values_out, +//! num_items, num_segments, d_offsets, d_offsets + 1); +//! +//! // d_keys_out <-- [6, 7, 8, 0, 3, 5, 9] +//! // d_values_out <-- [1, 2, 0, 5, 4, 3, 6] +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceSegmentedSort that accept an environment can be tuned by passing a custom :ref:`policy +//! selector ` that returns a :cpp:struct:`cub::SegmentedSortPolicy`, as shown in the example +//! below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin sort-keys-custom-policy-selector +//! :end-before: example-end sort-keys-custom-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin sort-keys-custom-policy +//! :end-before: example-end sort-keys-custom-policy +//! +//! @endrst +struct DeviceSegmentedSort +{ +private: + // Name reported for NVTX ranges + _CCCL_HOST_DEVICE static constexpr auto GetName() -> const char* + { + return "cub::DeviceSegmentedSort"; + } + + // TODO(bgruber): I would ideally like to have the logic of extracting the policy selector from the tuning environment + // inside the dispatch function, but this will not work with CCCL.C, which needs to pass a stateful policy selector. + // Refactor this once we have a host code JIT compiler. + template > + CUB_RUNTIME_FUNCTION static auto select_tuning_and_dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + detail::segmented_sort::global_segment_offset_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + bool is_overwrite_okay, + cudaStream_t stream, + TuningEnvT = {}) -> cudaError_t + { + using offset_t = + detail::choose_signed_offset_t>; + using default_policy_selector_t = detail::segmented_sort::policy_selector_from_types; + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + return detail::segmented_sort::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + is_overwrite_okay, + stream, + policy_selector_t{}); + } + + template > + CUB_RUNTIME_FUNCTION static cudaError_t sort_keys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream, + bool is_overwrite_okay = true, + TuningEnvT tuning_env = {}) + { + DoubleBuffer d_values; + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + is_overwrite_okay, + stream, + tuning_env); + } + + template > + CUB_RUNTIME_FUNCTION static cudaError_t sort_keys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream, + TuningEnvT tuning_env = {}) + { + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + false, + tuning_env); + } + +public: + //! @name Keys-only + //! @{ + + //! @rst + //! Sorts segments of keys into ascending order. + //! Approximately ``num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as `segment_offsets+1`). + //! - SortKeys is not guaranteed to be stable. That is, suppose that ``i`` and + //! ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``int`` keys. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible + //! // pointers for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [6, 7, 8, 0, 3, 5, 9] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the i-th segment is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of keys into ascending order. + //! Approximately ``num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - ``SortKeys`` is not guaranteed to be stable. That is, suppose that ``i`` and + //! ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-keys-env + //! :end-before: example-end sort-keys-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the i-th segment is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + tuning); + }); + } + + //! @rst + //! Sorts segments of keys into descending order. Approximately + //! ``num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortKeysDescending is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is + //! not guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [8, 7, 6, 9, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of keys into descending order. + //! Approximately ``num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - ``SortKeysDescending`` is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is + //! not guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-keys-descending-env + //! :end-before: example-end sort-keys-descending-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + tuning); + }); + } + + //! @rst + //! Sorts segments of keys into ascending order. Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! ``DoubleBuffer`` structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the ``DoubleBuffer`` wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets +1``). + //! - SortKeys is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is + //! not guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible + //! // pointers for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::SortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [6, 7, 8, 0, 3, 5, 9] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_keys( + d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, stream); + } + + //! @rst + //! Sorts segments of keys into ascending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! ``DoubleBuffer`` structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the ``DoubleBuffer`` wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - ``SortKeys`` is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is + //! not guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-keys-db-env + //! :end-before: example-end sort-keys-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeys( + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + true, + tuning); + }); + } + + //! @rst + //! Sorts segments of keys into descending order. Approximately + //! ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! ``DoubleBuffer`` structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the ``DoubleBuffer`` wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortKeysDescending is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is + //! not guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::SortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [8, 7, 6, 9, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1<= d_begin_offsets[i]``, the ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_keys( + d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, stream); + } + + //! @rst + //! Sorts segments of keys into descending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! ``DoubleBuffer`` structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the ``DoubleBuffer`` wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - ``SortKeysDescending`` is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is + //! not guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-keys-descending-db-env + //! :end-before: example-end sort-keys-descending-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1<= d_begin_offsets[i]``, the ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortKeysDescending( + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + true, + tuning); + }); + } + + //! @rst + //! Sorts segments of keys into ascending order. Approximately + //! ``num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortKeys is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::StableSortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::StableSortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [6, 7, 8, 0, 3, 5, 9] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of keys into ascending order. + //! Approximately ``num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - ``StableSortKeys`` is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-keys-env + //! :end-before: example-end stable-sort-keys-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeys( + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + tuning); + }); + } + + //! @rst + //! Sorts segments of keys into descending order. + //! Approximately ``num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortKeysDescending is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither ``x < y`` nor ``y < x``) + //! then a postcondition of stable sort is that ``x`` still precedes ``y``. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::StableSortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::StableSortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys_in, d_keys_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [8, 7, 6, 9, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and + //! ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of keys into descending order. + //! Approximately ``num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - ``StableSortKeysDescending`` is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither ``x < y`` nor ``y < x``) + //! then a postcondition of stable sort is that ``x`` still precedes ``y``. + //! - The range ``[d_keys_out, d_keys_out + num_items)`` shall not overlap + //! ``[d_keys_in, d_keys_in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_keys_out[i]`` will not + //! be accessed nor modified. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-keys-descending-env + //! :end-before: example-end stable-sort-keys-descending-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and + //! ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeysDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + tuning); + }); + } + + //! @rst + //! Sorts segments of keys into ascending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! ``DoubleBuffer`` structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the ``DoubleBuffer`` wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortKeys is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::StableSortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::StableSortKeys( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [6, 7, 8, 0, 3, 5, 9] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeys( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_keys( + d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, stream); + } + + //! @rst + //! Sorts segments of keys into ascending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! ``DoubleBuffer`` structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the ``DoubleBuffer`` wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - ``StableSortKeys`` is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-keys-db-env + //! :end-before: example-end stable-sort-keys-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeys( + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + true, + tuning); + }); + } + + //! @rst + //! Sorts segments of keys into descending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! ``DoubleBuffer`` structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the ``DoubleBuffer`` wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortKeysDescending is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a DoubleBuffer to wrap the pair of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::StableSortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::StableSortKeysDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [8, 7, 6, 9, 5, 3, 0] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and + //! ``d_values_*``. If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the + //! ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeysDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_keys( + d_temp_storage, temp_storage_bytes, d_keys, num_items, num_segments, d_begin_offsets, d_end_offsets, stream); + } + + //! @rst + //! Sorts segments of keys into descending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The sorting operation is given a pair of key buffers managed by a + //! ``DoubleBuffer`` structure that indicates which of the two buffers is + //! "current" (and thus contains the input data to be sorted). + //! - The contents of both buffers may be altered by the sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within the ``DoubleBuffer`` wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - ``StableSortKeysDescending`` is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``cur = d_keys.Current()`` and ``alt = d_keys.Alternate()``. + //! The range ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_keys[i].Alternate()[i]`` will not be accessed nor modified. + //! - Can use a specific stream or cuda memory resource through the ``env`` parameter. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_keys_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-keys-descending-db-env + //! :end-before: example-end stable-sort-keys-descending-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and + //! ``d_values_*``. If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the + //! ``i``-th segment is considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t StableSortKeysDescending( + DoubleBuffer& d_keys, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_keys( + d_temp_storage, + temp_storage_bytes, + d_keys, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + true, + tuning); + }); + } + +private: + template > + CUB_RUNTIME_FUNCTION static cudaError_t sort_pairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream, + bool is_overwrite_okay = true, + TuningEnvT tuning_env = {}) + { + return select_tuning_and_dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + is_overwrite_okay, + stream, + tuning_env); + } + + template > + CUB_RUNTIME_FUNCTION static cudaError_t sort_pairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream, + TuningEnvT tuning_env = {}) + { + DoubleBuffer d_keys(const_cast(d_keys_in), d_keys_out); + DoubleBuffer d_values(const_cast(d_values_in), d_values_out); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + false, + tuning_env); + } + +public: + //! @} + //! @name Key-value pairs + //! @{ + + //! @rst + //! Sorts segments of key-value pairs into ascending order. + //! Approximately ``2 * num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortPairs is not guaranteed to be stable. That is, suppose that ``i`` and + //! ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys with associated vector of + //! ``i`` nt values. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_values_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [6, 7, 8, 0, 3, 5, 9] + //! // d_values_out <-- [1, 2, 0, 5, 4, 3, 6] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i]-1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into ascending order. + //! Approximately ``2 * num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortPairs is not guaranteed to be stable. That is, suppose that ``i`` and + //! ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_pairs_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-pairs-env + //! :end-before: example-end sort-pairs-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + tuning); + }); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. + //! Approximately ``2 * num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortPairsDescending is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys with associated vector of + //! ``i`` nt values. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_values_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [8, 7, 6, 9, 5, 3, 0] + //! // d_values_out <-- [0, 2, 1, 6, 3, 4, 5] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the i-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. + //! Approximately ``2 * num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortPairsDescending is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_pairs_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-pairs-descending-env + //! :end-before: example-end sort-pairs-descending-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + tuning); + }); + } + + //! @rst + //! Sorts segments of key-value pairs into ascending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting + //! operation. + //! - Upon completion, the sorting operation will update the "current" indicator + //! within each DoubleBuffer wrapper to reference which of the two buffers + //! now contains the sorted output sequence (a function of the number of key bits + //! specified and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortPairs is not guaranteed to be stable. That is, suppose that ``i`` and + //! ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys with associated vector of + //! ``i`` nt values. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! int *d_value_buf; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_value_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a set of DoubleBuffers to wrap pairs of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! cub::DoubleBuffer d_values(d_value_buf, d_value_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::SortPairs( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [6, 7, 8, 0, 3, 5, 9] + //! // d_values.Current() <-- [5, 4, 3, 1, 2, 0, 6] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer contains + //! the unsorted input values and, upon return, is updated to point to the + //! sorted output values + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the i-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into ascending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting + //! operation. + //! - Upon completion, the sorting operation will update the "current" indicator + //! within each DoubleBuffer wrapper to reference which of the two buffers + //! now contains the sorted output sequence (a function of the number of key bits + //! specified and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortPairs is not guaranteed to be stable. That is, suppose that ``i`` and + //! ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_pairs_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-pairs-db-env + //! :end-before: example-end sort-pairs-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer contains + //! the unsorted input values and, upon return, is updated to point to the + //! sorted output values + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortPairs( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + true, + tuning); + }); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits specified and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortPairsDescending is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is + //! not guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys with associated vector of + //! ``i`` nt values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for + //! // sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! int *d_value_buf; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_value_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a set of DoubleBuffers to wrap pairs of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! cub::DoubleBuffer d_values(d_value_buf, d_value_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::SortPairsDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [8, 7, 6, 9, 5, 3, 0] + //! // d_values.Current() <-- [0, 2, 1, 6, 3, 4, 5] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer contains + //! the unsorted input values and, upon return, is updated to point to the + //! sorted output values + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting + //! operation. + //! - Upon completion, the sorting operation will update the "current" indicator + //! within each DoubleBuffer wrapper to reference which of the two buffers + //! now contains the sorted output sequence (a function of the number of key bits + //! specified and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - SortPairsDescending is not guaranteed to be stable. That is, suppose that + //! ``i`` and ``j`` are equivalent: neither one is less than the other. It is not + //! guaranteed that the relative order of these two elements will be + //! preserved by sort. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_pairs_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin sort-pairs-descending-db-env + //! :end-before: example-end sort-pairs-descending-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer contains + //! the unsorted input values and, upon return, is updated to point to the + //! sorted output values + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t SortPairsDescending( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + true, + tuning); + }); + } + + //! @rst + //! Sorts segments of key-value pairs into ascending order. + //! Approximately ``2 * num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortPairs is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys with associated vector of + //! ``i`` nt values. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_values_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::StableSortPairs( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::StableSortPairs( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [6, 7, 8, 0, 3, 5, 9] + //! // d_values_out <-- [1, 2, 0, 5, 4, 3, 6] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into ascending order. + //! Approximately ``2 * num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortPairs is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! + //! Snippet + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_pairs_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-pairs-env + //! :end-before: example-end stable-sort-pairs-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairs( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + tuning); + }); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. + //! Approximately ``2 * num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortPairsDescending is stable: it preserves the relative ordering + //! of equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let `in` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys with associated vector of + //! ``i`` nt values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_keys_in; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_keys_out; // e.g., [-, -, -, -, -, -, -] + //! int *d_values_in; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_values_out; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::StableSortPairsDescending( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::StableSortPairsDescending( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_keys_out, d_values_in, d_values_out, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys_out <-- [8, 7, 6, 9, 5, 3, 0] + //! // d_values_out <-- [0, 2, 1, 6, 3, 4, 5] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @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 + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. + //! Approximately ``2 * num_items + 2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The contents of the input data are not altered by the sorting operation. + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortPairsDescending is stable: it preserves the relative ordering + //! of equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``in`` be one of ``{d_keys_in, d_values_in}`` and ``out`` be any of + //! ``{d_keys_out, d_values_out}``. The range ``[out, out + num_items)`` shall + //! not overlap ``[in, in + num_items)``, + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys_in[i]``, ``d_values_in[i]``, + //! ``d_keys_out[i]``, ``d_values_out[i]`` will not be accessed nor modified. + //! + //! Snippet + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_pairs_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-pairs-descending-env + //! :end-before: example-end stable-sort-pairs-descending-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in] d_keys_in + //! Device-accessible pointer to the input data of key data to sort + //! + //! @param[out] d_keys_out + //! Device-accessible pointer to the sorted output sequence of key data + //! + //! @param[in] d_values_in + //! Device-accessible pointer to the corresponding input sequence of + //! associated value items + //! + //! @param[out] d_values_out + //! Device-accessible pointer to the correspondingly-reordered output + //! sequence of associated value items + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairsDescending( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + tuning); + }); + } + + //! @rst + //! Sorts segments of key-value pairs into ascending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the + //! sorting operation. + //! - Upon completion, the sorting operation will update the "current" + //! indicator within each DoubleBuffer wrapper to reference which of the two + //! buffers now contains the sorted output sequence (a function of the number + //! of key bits specified and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortPairs is stable: it preserves the relative ordering + //! of equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes `y`, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys with associated vector of + //! ``i`` nt values. + //! + //! .. code-block:: c++ + //! + //! #include + //! // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! int *d_value_buf; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_value_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a set of DoubleBuffers to wrap pairs of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! cub::DoubleBuffer d_values(d_value_buf, d_value_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::StableSortPairs( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::StableSortPairs( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [6, 7, 8, 0, 3, 5, 9] + //! // d_values.Current() <-- [5, 4, 3, 1, 2, 0, 6] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer contains + //! the unsorted input values and, upon return, is updated to point to the + //! sorted output values + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i]-1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairs( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into ascending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting + //! operation. + //! - Upon completion, the sorting operation will update the "current" indicator + //! within each DoubleBuffer wrapper to reference which of the two buffers + //! now contains the sorted output sequence (a function of the number of key bits + //! specified and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortPairs is stable: it preserves the relative ordering of + //! equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_pairs_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-pairs-db-env + //! :end-before: example-end stable-sort-pairs-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer contains + //! the unsorted input values and, upon return, is updated to point to the + //! sorted output values + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairs( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + true, + tuning); + }); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting + //! operation. + //! - Upon completion, the sorting operation will update the "current" indicator + //! within each DoubleBuffer wrapper to reference which of the two buffers + //! now contains the sorted output sequence (a function of the number of key bits + //! specified and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortPairsDescending is stable: it preserves the relative ordering + //! of equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the batched sorting of three segments + //! (with one zero-length segment) of ``i`` nt keys with associated vector of + //! ``i`` nt values. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for sorting data + //! int num_items; // e.g., 7 + //! int num_segments; // e.g., 3 + //! int *d_offsets; // e.g., [0, 3, 3, 7] + //! int *d_key_buf; // e.g., [8, 6, 7, 5, 3, 0, 9] + //! int *d_key_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! int *d_value_buf; // e.g., [0, 1, 2, 3, 4, 5, 6] + //! int *d_value_alt_buf; // e.g., [-, -, -, -, -, -, -] + //! ... + //! + //! // Create a set of DoubleBuffers to wrap pairs of device pointers + //! cub::DoubleBuffer d_keys(d_key_buf, d_key_alt_buf); + //! cub::DoubleBuffer d_values(d_value_buf, d_value_alt_buf); + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSegmentedSort::StableSortPairsDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run sorting operation + //! cub::DeviceSegmentedSort::StableSortPairsDescending( + //! d_temp_storage, temp_storage_bytes, d_keys, d_values, + //! num_items, num_segments, d_offsets, d_offsets + 1); + //! + //! // d_keys.Current() <-- [8, 7, 6, 9, 5, 3, 0] + //! // d_values.Current() <-- [0, 2, 1, 6, 3, 4, 5] + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer contains + //! the unsorted input values and, upon return, is updated to point to the + //! sorted output values + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] stream + //! @rst + //! **[optional]** CUDA stream to launch kernels within. Default is stream\ :sub:`0`. + //! @endrst + template + CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairsDescending( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream = nullptr) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, GetName()); + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream); + } + + //! @rst + //! Sorts segments of key-value pairs into descending order. + //! Approximately ``2 * num_segments`` auxiliary storage required. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The sorting operation is given a pair of key buffers and a corresponding + //! pair of associated value buffers. Each pair is managed by a DoubleBuffer + //! structure that indicates which of the two buffers is "current" (and thus + //! contains the input data to be sorted). + //! - The contents of both buffers within each pair may be altered by the sorting + //! operation. + //! - Upon completion, the sorting operation will update the "current" indicator + //! within each DoubleBuffer wrapper to reference which of the two buffers + //! now contains the sorted output sequence (a function of the number of key bits + //! specified and the targeted device architecture). + //! - When the input is a contiguous sequence of segments, a single sequence + //! ``segment_offsets`` (of length ``num_segments + 1``) can be aliased + //! for both the ``d_begin_offsets`` and ``d_end_offsets`` parameters (where + //! the latter is specified as ``segment_offsets + 1``). + //! - StableSortPairsDescending is stable: it preserves the relative ordering + //! of equivalent elements. That is, if ``x`` and ``y`` are elements such that + //! ``x`` precedes ``y``, and if the two elements are equivalent (neither + //! ``x < y`` nor ``y < x``) then a postcondition of stable sort is that + //! ``x`` still precedes ``y``. + //! - Let ``cur`` be one of ``{d_keys.Current(), d_values.Current()}`` and ``alt`` + //! be any of ``{d_keys.Alternate(), d_values.Alternate()}``. The range + //! ``[cur, cur + num_items)`` shall not overlap + //! ``[alt, alt + num_items)``. Both ranges shall not overlap + //! ``[d_begin_offsets, d_begin_offsets + num_segments)`` nor + //! ``[d_end_offsets, d_end_offsets + num_segments)`` in any way. + //! - Segments are not required to be contiguous. For all index values ``i`` + //! outside the specified segments ``d_keys.Current()[i]``, + //! ``d_values.Current()[i]``, ``d_keys.Alternate()[i]``, + //! ``d_values.Alternate()[i]`` will not be accessed nor modified. + //! + //! Snippet + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_segmented_sort_pairs_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin stable-sort-pairs-descending-db-env + //! :end-before: example-end stable-sort-pairs-descending-db-env + //! + //! @endrst + //! + //! @tparam KeyT + //! **[inferred]** Key type + //! + //! @tparam ValueT + //! **[inferred]** Value type + //! + //! @tparam BeginOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! beginning offsets @iterator + //! + //! @tparam EndOffsetIteratorT + //! **[inferred]** Random-access input iterator type for reading segment + //! ending offsets @iterator + //! + //! @tparam EnvT + //! **[optional]** Environment type providing stream and other properties + //! + //! @param[in,out] d_keys + //! Reference to the double-buffer of keys whose "current" device-accessible + //! buffer contains the unsorted input keys and, upon return, is updated to + //! point to the sorted output keys + //! + //! @param[in,out] d_values + //! Double-buffer of values whose "current" device-accessible buffer contains + //! the unsorted input values and, upon return, is updated to point to the + //! sorted output values + //! + //! @param[in] num_items + //! The total number of items to sort (across all segments) + //! + //! @param[in] num_segments + //! The number of segments that comprise the sorting data + //! + //! @param[in] d_begin_offsets + //! @rst + //! Random-access input iterator to the sequence of beginning offsets of + //! length ``num_segments``, such that ``d_begin_offsets[i]`` is the first + //! element of the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*`` + //! @endrst + //! + //! @param[in] d_end_offsets + //! @rst + //! Random-access input iterator to the sequence of ending offsets of length + //! ``num_segments``, such that ``d_end_offsets[i] - 1`` is the last element of + //! the *i*\ :sup:`th` data segment in ``d_keys_*`` and ``d_values_*``. + //! If ``d_end_offsets[i] - 1 <= d_begin_offsets[i]``, the ``i``-th segment is + //! considered empty. + //! @endrst + //! + //! @param[in] env + //! @rst + //! **[optional]** Environment providing stream and other properties. Default is empty environment. + //! @endrst + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION static cudaError_t StableSortPairsDescending( + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE(GetName()); + return detail::dispatch_with_env( + env, [&]([[maybe_unused]] auto tuning, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return sort_pairs( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + true, + tuning); + }); + } + + //! @} +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_select.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_select.cuh new file mode 100644 index 00000000..44b8d636 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_select.cuh @@ -0,0 +1,2805 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! cub::DeviceSelect provides device-wide, parallel operations for compacting selected items from sequences of data +//! items residing within device-accessible memory. + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +//! @rst +//! DeviceSelect provides device-wide, parallel operations for compacting selected items from sequences of data items +//! residing within device-accessible memory. It is similar to DevicePartition, except that non-selected items are +//! discarded, whereas DevicePartition retains them. +//! +//! Overview +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! These operations apply a selection criterion to selectively copy +//! items from a specified input sequence to a compact output sequence. +//! +//! Usage Considerations +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @cdp_class{DeviceSelect} +//! +//! Performance +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! @linear_performance{select-flagged, select-if, and select-unique} +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All non-ByKey algorithms in DeviceSelect that accept an environment can be tuned by passing a custom +//! :ref:`policy selector ` that returns a :cpp:struct:`cub::SelectPolicy`, as shown in the +//! example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin select-if-policy-selector +//! :end-before: example-end select-if-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin select-if-tuning +//! :end-before: example-end select-if-tuning +//! +//! All *ByKey algorithms in DeviceSelect that accept an environment can be tuned by passing a +//! custom :ref:`policy selector ` that returns a :cpp:struct:`cub::UniqueByKeyPolicy`, as shown +//! in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin unique-by-key-policy-selector +//! :end-before: example-end unique-by-key-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin unique-by-key-tuning +//! :end-before: example-end unique-by-key-tuning +//! +//! @endrst +struct DeviceSelect +{ + //! @rst + //! Uses the ``d_flags`` sequence to selectively copy the corresponding items from ``d_in`` into ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The value type of ``d_flags`` must be castable to ``bool`` (e.g., ``bool``, ``char``, ``int``, etc.). + //! - Copies of the selected items are compacted into ``d_out`` and maintain their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap ``[d_in, d_in + num_items)``, + //! | ``[d_flags, d_flags + num_items)`` nor ``d_num_selected_out`` in any way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for input, + //! // flags, and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [1, 2, 3, 4, 5, 6, 7, 8] + //! char *d_flags; // e.g., [1, 0, 0, 1, 0, 1, 1, 0] + //! int *d_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSelect::Flagged( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_flags, d_out, d_num_selected_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DeviceSelect::Flagged( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_flags, d_out, d_num_selected_out, num_items); + //! + //! // d_out <-- [1, 4, 6, 7] + //! // d_num_selected_out <-- [4] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Flagged( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FlagIterator d_flags, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::Flagged"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + d_flags, + d_out, + d_num_selected_out, + NullType{}, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``d_flags`` sequence to selectively copy the corresponding items from ``d_in`` into ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The value type of ``d_flags`` must be castable to ``bool`` (e.g., ``bool``, ``char``, ``int``, etc.). + //! - Copies of the selected items are compacted into ``d_out`` and maintain their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap ``[d_in, d_in + num_items)``, + //! | ``[d_flags, d_flags + num_items)`` nor ``d_num_selected_out`` in any way. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of flagged items from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-flagged-env + //! :end-before: example-end select-flagged-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Flagged( + InputIteratorT d_in, + FlagIterator d_flags, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::Flagged"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + d_flags, + d_out, + d_num_selected_out, + NullType{}, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``d_flags`` sequence to selectively compact items in ``d_data``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The value type of ``d_flags`` must be castable to ``bool`` (e.g., ``bool``, ``char``, ``int``, etc.). + //! - Copies of the selected items are compacted in-place and maintain their original relative ordering. + //! - | The ``d_data`` may equal ``d_flags``. The range ``[d_data, d_data + num_items)`` shall not overlap + //! | ``[d_flags, d_flags + num_items)`` in any other way. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the in-place compaction of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-flagged-inplace-env + //! :end-before: example-end select-flagged-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing selected items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + Flagged(IteratorT d_data, + FlagIterator d_flags, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::Flagged"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + FlagIterator, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + d_flags, + d_data, + d_num_selected_out, + NullType{}, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``select_op`` functor to selectively copy items from ``d_in`` into ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap + //! | ``[d_in, d_in + num_items)`` nor ``d_num_selected_out`` in any way. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-if-env + //! :end-before: example-end select-if-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection operator type having member `bool operator()(const T &a)` + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_unary_predicate, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + If(InputIteratorT d_in, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::If"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + static_cast(nullptr), + d_out, + d_num_selected_out, + select_op, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``select_op`` functor to selectively compact items in ``d_data``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - Copies of the selected items are compacted in ``d_data`` and maintain + //! their original relative ordering. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the in-place compaction of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-if-inplace-env + //! :end-before: example-end select-if-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection operator type having member `bool operator()(const T &a)` + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_unary_predicate, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + If(IteratorT d_data, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::If"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + NullType*, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + static_cast(nullptr), + d_data, + d_num_selected_out, + select_op, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``d_flags`` sequence to selectively compact the items in `d_data``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The value type of ``d_flags`` must be castable to ``bool`` (e.g., ``bool``, ``char``, ``int``, etc.). + //! - Copies of the selected items are compacted in-place and maintain their original relative ordering. + //! - | The ``d_data`` may equal ``d_flags``. The range ``[d_data, d_data + num_items)`` shall not overlap + //! | ``[d_flags, d_flags + num_items)`` in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers for input, + //! // flags, and output + //! int num_items; // e.g., 8 + //! int *d_data; // e.g., [1, 2, 3, 4, 5, 6, 7, 8] + //! char *d_flags; // e.g., [1, 0, 0, 1, 0, 1, 1, 0] + //! int *d_num_selected_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSelect::Flagged( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_flags, d_num_selected_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DeviceSelect::Flagged( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_flags, d_num_selected_out, num_items); + //! + //! // d_data <-- [1, 4, 6, 7] + //! // d_num_selected_out <-- [4] + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing selected items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Flagged( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + FlagIterator d_flags, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::Flagged"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + FlagIterator, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + d_flags, + d_data, + d_num_selected_out, + NullType{}, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``select_op`` functor to selectively copy items from ``d_in`` into ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap + //! | ``[d_in, d_in + num_items)`` nor ``d_num_selected_out`` in any way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Functor type for selecting values less than some criteria + //! struct LessThan + //! { + //! int compare; + //! + //! __host__ __device__ __forceinline__ + //! LessThan(int compare) : compare(compare) {} + //! + //! __host__ __device__ __forceinline__ + //! bool operator()(const int &a) const { + //! return (a < compare); + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [0, 2, 3, 9, 5, 2, 81, 8] + //! int *d_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ ] + //! LessThan select_op(7); + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSelect::If( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, d_num_selected_out, num_items, select_op); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DeviceSelect::If( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, d_num_selected_out, num_items, select_op); + //! + //! // d_out <-- [0, 2, 3, 5, 2] + //! // d_num_selected_out <-- [5] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection operator type having member `bool operator()(const T &a)` + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_unary_predicate, int> = 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + If(void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::If"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + static_cast(nullptr), + d_out, + d_num_selected_out, + select_op, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``select_op`` functor to selectively compact items in ``d_data``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - | Copies of the selected items are compacted in ``d_data`` and maintain + //! | their original relative ordering. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Functor type for selecting values less than some criteria + //! struct LessThan + //! { + //! int compare; + //! + //! __host__ __device__ __forceinline__ + //! LessThan(int compare) : compare(compare) {} + //! + //! __host__ __device__ __forceinline__ + //! bool operator()(const int &a) const { + //! return (a < compare); + //! } + //! }; + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 8 + //! int *d_data; // e.g., [0, 2, 3, 9, 5, 2, 81, 8] + //! int *d_num_selected_out; // e.g., [ ] + //! LessThan select_op(7); + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSelect::If( + //! d_temp_storage, temp_storage_bytes, + //! d_data, d_num_selected_out, num_items, select_op); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DeviceSelect::If( + //! d_temp_storage, temp_storage_bytes, + //! d_data, d_num_selected_out, num_items, select_op); + //! + //! // d_data <-- [0, 2, 3, 5, 2] + //! // d_num_selected_out <-- [5] + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access input iterator type for reading and writing items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection operator type having member `bool operator()(const T &a)` + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_unary_predicate, int> = 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + If(void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::If"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + NullType*, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + static_cast(nullptr), + d_data, + d_num_selected_out, + select_op, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``select_op`` functor applied to ``d_flags`` to selectively copy the + //! corresponding items from ``d_in`` into ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The expression ``select_op(flag)`` must be convertible to ``bool``, + //! where the type of ``flag`` corresponds to the value type of ``FlagIterator``. + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap + //! | ``[d_in, d_in + num_items)`` nor ``d_num_selected_out`` in any way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-select-iseven + //! :end-before: example-end segmented-select-iseven + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-select-flaggedif + //! :end-before: example-end segmented-select-flaggedif + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection operator type having member `bool operator()(const T &a)` + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t FlaggedIf( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FlagIterator d_flags, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::FlaggedIf"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + d_flags, + d_out, + d_num_selected_out, + select_op, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``select_op`` functor applied to ``d_flags`` to selectively compact the + //! corresponding items in ``d_data``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The expression ``select_op(flag)`` must be convertible to ``bool``, + //! where the type of ``flag`` corresponds to the value type of ``FlagIterator``. + //! - Copies of the selected items are compacted in-place and maintain their original relative ordering. + //! - | The ``d_data`` may equal ``d_flags``. The range ``[d_data, d_data + num_items)`` shall not overlap + //! | ``[d_flags, d_flags + num_items)`` in any other way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-select-iseven + //! :end-before: example-end segmented-select-iseven + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin segmented-select-flaggedif-inplace + //! :end-before: example-end segmented-select-flaggedif-inplace + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing selected items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection operator type having member `bool operator()(const T &a)` + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t FlaggedIf( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + FlagIterator d_flags, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::FlaggedIf"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + FlagIterator, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + d_flags, + d_data, + d_num_selected_out, + select_op, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``select_op`` functor applied to ``d_flags`` to selectively copy the + //! corresponding items from ``d_in`` into ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The expression ``select_op(flag)`` must be convertible to ``bool``, + //! where the type of ``flag`` corresponds to the value type of ``FlagIterator``. + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap + //! | ``[d_in, d_in + num_items)`` nor ``d_num_selected_out`` in any way. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-flaggedif-env + //! :end-before: example-end select-flaggedif-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection operator type having member `bool operator()(const T &a)` + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_unary_predicate, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t FlaggedIf( + InputIteratorT d_in, + FlagIterator d_flags, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::FlaggedIf"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + d_flags, + d_out, + d_num_selected_out, + select_op, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Uses the ``select_op`` functor applied to ``d_flags`` to selectively compact + //! items in ``d_data``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The expression ``select_op(flag)`` must be convertible to ``bool``, + //! where the type of ``flag`` corresponds to the value type of ``FlagIterator``. + //! - Copies of the selected items are compacted in-place and maintain their original relative ordering. + //! - | The ``d_data`` may equal ``d_flags``. The range ``[d_data, d_data + num_items)`` shall not overlap + //! | ``[d_flags, d_flags + num_items)`` in any other way. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the in-place compaction of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-flaggedif-inplace-env + //! :end-before: example-end select-flaggedif-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing selected items @iterator + //! + //! @tparam FlagIterator + //! **[inferred]** Random-access input iterator type for reading selection flags @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam SelectOp + //! **[inferred]** Selection operator type having member `bool operator()(const T &a)` + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[in] d_flags + //! Pointer to the input sequence of selection flags + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] select_op + //! Unary selection operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_unary_predicate, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t FlaggedIf( + IteratorT d_data, + FlagIterator d_flags, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + SelectOp select_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::FlaggedIf"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + FlagIterator, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + d_flags, + d_data, + d_num_selected_out, + select_op, + NullType{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_in`` having runs of consecutive equal-valued keys, + //! only the first key from each run is selectively copied to ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The ``==`` equality operator is used to determine whether keys are equivalent. + //! - Copies of the selected items are compacted into ``d_out`` and maintain their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap + //! | ``[d_in, d_in + num_items)`` nor ``d_num_selected_out`` in any way. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-unique-env + //! :end-before: example-end select-unique-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template < + typename InputIteratorT, + typename OutputIteratorT, + typename NumSelectedIteratorT, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + Unique(InputIteratorT d_in, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::Unique"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + static_cast(nullptr), + d_out, + d_num_selected_out, + NullType{}, + ::cuda::std::equal_to<>{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_in`` having runs of consecutive equal-valued keys, + //! only the first key from each run is selectively copied to ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The user-provided equality operator, ``equality_op``, is used to determine whether keys are equivalent. + //! - Copies of the selected items are compacted into ``d_out`` and maintain their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap + //! | ``[d_in, d_in + num_items)`` nor ``d_num_selected_out`` in any way. + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector + //! using a custom equality operator: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-unique-eqop-env + //! :end-before: example-end select-unique-eqop-env + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EqualityOpT + //! **[inferred]** Type of equality_op + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in] d_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] equality_op + //! Binary equality operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_binary_predicate, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Unique( + InputIteratorT d_in, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + EqualityOpT equality_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::Unique"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + static_cast(nullptr), + d_out, + d_num_selected_out, + NullType{}, + equality_op, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_data`` having runs of consecutive equal-valued keys, + //! only the first key from each run is selectively compacted in-place. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The ``==`` equality operator is used to determine whether keys are equivalent. + //! - Copies of the selected items are compacted in ``d_data`` and maintain their original relative ordering. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the in-place compaction of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-unique-inplace-env + //! :end-before: example-end select-unique-inplace-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Unique( + IteratorT d_data, NumSelectedIteratorT d_num_selected_out, ::cuda::std::int64_t num_items, const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::Unique"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + NullType*, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + static_cast(nullptr), + d_data, + d_num_selected_out, + NullType{}, + ::cuda::std::equal_to<>{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_data`` having runs of consecutive equal-valued keys, + //! only the first key from each run is selectively compacted in-place. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The user-provided equality operator, ``equality_op``, is used to determine whether keys are equivalent. + //! - Copies of the selected items are compacted in ``d_data`` and maintain their original relative ordering. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the in-place compaction of items selected from an ``int`` device vector + //! using a custom equality operator: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-unique-inplace-eqop-env + //! :end-before: example-end select-unique-inplace-eqop-env + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EqualityOpT + //! **[inferred]** Type of equality_op + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., ``cuda::std::execution::env<...>``) + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] equality_op + //! Binary equality operator + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_binary_predicate, int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t + Unique(IteratorT d_data, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + EqualityOpT equality_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::Unique"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + NullType*, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + static_cast(nullptr), + d_data, + d_num_selected_out, + NullType{}, + equality_op, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_keys_in`` and ``d_values_in`` with runs of key-value pairs with consecutive + //! equal-valued keys, only the first key and its value from each run is selectively copied + //! to ``d_keys_out`` and ``d_values_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The user-provided equality operator, ``equality_op``, is used to determine whether keys are equivalent. + //! - Copies of the selected items are compacted into ``d_keys_out`` and ``d_values_out`` and maintain + //! their original relative ordering. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector + //! using environment-based API: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-uniquebykey-env + //! :end-before: example-end select-uniquebykey-env + //! + //! @tparam KeyInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input keys @iterator + //! + //! @tparam ValueInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input values @iterator + //! + //! @tparam KeyOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected keys @iterator + //! + //! @tparam ValueOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected values @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EqualityOpT + //! **[inferred]** Type of equality_op + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_keys_in + //! Pointer to the input sequence of keys + //! + //! @param[in] d_values_in + //! Pointer to the input sequence of values + //! + //! @param[out] d_keys_out + //! Pointer to the output sequence of selected keys + //! + //! @param[out] d_values_out + //! Pointer to the output sequence of selected values + //! + //! @param[out] d_num_selected_out + //! Pointer to the total number of items selected (i.e., length of `d_keys_out` or `d_values_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_keys_in` or `d_values_in`) + //! + //! @param[in] equality_op + //! Binary predicate to determine equality + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template < + typename KeyInputIteratorT, + typename ValueInputIteratorT, + typename KeyOutputIteratorT, + typename ValueOutputIteratorT, + typename NumSelectedIteratorT, + typename NumItemsT, + typename EqualityOpT, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0, + ::cuda::std::enable_if_t<::cuda::std::indirect_binary_predicate, + int> = 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t UniqueByKey( + KeyInputIteratorT d_keys_in, + ValueInputIteratorT d_values_in, + KeyOutputIteratorT d_keys_out, + ValueOutputIteratorT d_values_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + EqualityOpT equality_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceSelect::UniqueByKey"); + + using offset_t = detail::choose_offset_t; + + using default_policy_selector = + detail::unique_by_key::policy_selector_from_types, + detail::it_value_t>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::unique_by_key::dispatch( + storage, + bytes, + d_keys_in, + d_values_in, + d_keys_out, + d_values_out, + d_num_selected_out, + equality_op, + static_cast(num_items), + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_keys_in`` and ``d_values_in`` with runs of key-value pairs with consecutive + //! equal-valued keys, only the first key and its value from each run is selectively copied + //! to ``d_keys_out`` and ``d_values_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! 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`` + //! + //! - The ``==`` equality operator is used to determine whether keys are equivalent. + //! - Copies of the selected items are compacted into ``d_keys_out`` and ``d_values_out`` and maintain + //! their original relative ordering. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges. + //! + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector + //! using environment-based API and default key equality: + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_env_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-uniquebykey-default-eq-env + //! :end-before: example-end select-uniquebykey-default-eq-env + //! + //! @endrst + //! + //! @tparam KeyInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input keys @iterator + //! + //! @tparam ValueInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input values @iterator + //! + //! @tparam KeyOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected keys @iterator + //! + //! @tparam ValueOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected values @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_keys_in + //! Pointer to the input sequence of keys + //! + //! @param[in] d_values_in + //! Pointer to the input sequence of values + //! + //! @param[out] d_keys_out + //! Pointer to the output sequence of selected keys + //! + //! @param[out] d_values_out + //! Pointer to the output sequence of selected values + //! + //! @param[out] d_num_selected_out + //! Pointer to the total number of items selected (i.e., length of `d_keys_out` or `d_values_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_keys_in` or `d_values_in`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template < + typename KeyInputIteratorT, + typename ValueInputIteratorT, + typename KeyOutputIteratorT, + typename ValueOutputIteratorT, + typename NumSelectedIteratorT, + typename NumItemsT, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0, + ::cuda::std::enable_if_t, int> = + 0> + [[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t UniqueByKey( + KeyInputIteratorT d_keys_in, + ValueInputIteratorT d_values_in, + KeyOutputIteratorT d_keys_out, + ValueOutputIteratorT d_values_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + const EnvT& env = {}) + { + return UniqueByKey( + d_keys_in, d_values_in, d_keys_out, d_values_out, d_num_selected_out, num_items, ::cuda::std::equal_to<>{}, env); + } + + //! @rst + //! Given an input sequence ``d_in`` having runs of consecutive equal-valued keys, + //! only the first key from each run is selectively copied to ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The user-provided equality operator, `equality_op`, is used to determine whether keys are equivalent + //! - Copies of the selected items are compacted into ``d_out`` and maintain their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap + //! | ``[d_in, d_in + num_items)`` nor ``d_num_selected_out`` in any way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [0, 2, 2, 9, 5, 5, 5, 8] + //! int *d_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSelect::Unique( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, d_num_selected_out, num_items, cuda::std::equal_to<>{}); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DeviceSelect::Unique( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, d_num_selected_out, num_items, cuda::std::equal_to<>{}); + //! + //! // d_out <-- [0, 2, 9, 5, 8] + //! // d_num_selected_out <-- [5] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EqualityOpT + //! **[inferred]** Type of equality_op + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] equality_op + //! Binary predicate to determine equality + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_binary_predicate, + int> = 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Unique( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + EqualityOpT equality_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::Unique"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + static_cast(nullptr), + d_out, + d_num_selected_out, + NullType{}, + equality_op, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_in`` having runs of consecutive equal-valued keys, + //! only the first key from each run is selectively copied to ``d_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The ``==`` equality operator is used to determine whether keys are equivalent + //! - Copies of the selected items are compacted into ``d_out`` and maintain their original relative ordering. + //! - | The range ``[d_out, d_out + *d_num_selected_out)`` shall not overlap + //! | ``[d_in, d_in + num_items)`` nor ``d_num_selected_out`` in any way. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 8 + //! int *d_in; // e.g., [0, 2, 2, 9, 5, 5, 5, 8] + //! int *d_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSelect::Unique( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, d_num_selected_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DeviceSelect::Unique( + //! d_temp_storage, temp_storage_bytes, + //! d_in, d_out, d_num_selected_out, num_items); + //! + //! // d_out <-- [0, 2, 9, 5, 8] + //! // d_num_selected_out <-- [5] + //! + //! @endrst + //! + //! @tparam InputIteratorT + //! **[inferred]** Random-access input iterator type for reading input items @iterator + //! + //! @tparam OutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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_in + //! Pointer to the input sequence of data items + //! + //! @param[out] d_out + //! Pointer to the output sequence of selected data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! (i.e., length of `d_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_in`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template < + typename InputIteratorT, + typename OutputIteratorT, + typename NumSelectedIteratorT, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Unique( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::Unique"); + + using default_policy_selector = detail::select:: + policy_selector_from_types; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_in, + static_cast(nullptr), + d_out, + d_num_selected_out, + NullType{}, + ::cuda::std::equal_to<>{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_data`` having runs of consecutive equal-valued keys, + //! only the first key from each run is selectively compacted in-place. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The ``==`` equality operator is used to determine whether keys are equivalent + //! - Copies of the selected items are compacted in ``d_data`` and maintain their original relative ordering. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the in-place compaction of items selected from an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-unique-inplace + //! :end-before: example-end select-unique-inplace + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t, int> = 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Unique( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::Unique"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + NullType*, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + static_cast(nullptr), + d_data, + d_num_selected_out, + NullType{}, + ::cuda::std::equal_to<>{}, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_data`` having runs of consecutive equal-valued keys, + //! only the first key from each run is selectively compacted in-place. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 3.4.0 + //! First appears in CUDA Toolkit 13.4. + //! + //! - The user-provided equality operator, ``equality_op``, is used to determine whether keys are equivalent. + //! - Copies of the selected items are compacted in ``d_data`` and maintain their original relative ordering. + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the in-place compaction of items selected from an ``int`` device vector. + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-unique-inplace-eqop-myequalityop + //! :end-before: example-end select-unique-inplace-eqop-myequalityop + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_select_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin select-unique-inplace-eqop + //! :end-before: example-end select-unique-inplace-eqop + //! + //! @endrst + //! + //! @tparam IteratorT + //! **[inferred]** Random-access iterator type for reading and writing items @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam EqualityOpT + //! **[inferred]** Type of equality_op + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @param[in] d_temp_storage + //! @devicestorage + //! + //! @param[in,out] temp_storage_bytes + //! Reference to size in bytes of `d_temp_storage` allocation + //! + //! @param[in,out] d_data + //! Pointer to the sequence of data items + //! + //! @param[out] d_num_selected_out + //! Pointer to the output total number of items selected + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_data`) + //! + //! @param[in] equality_op + //! Binary predicate to determine equality + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template , + ::cuda::std::enable_if_t<::cuda::std::indirect_binary_predicate, int> = 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Unique( + void* d_temp_storage, + size_t& temp_storage_bytes, + IteratorT d_data, + NumSelectedIteratorT d_num_selected_out, + ::cuda::std::int64_t num_items, + EqualityOpT equality_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::Unique"); + + using default_policy_selector = detail::select::policy_selector_from_types< + IteratorT, + NullType*, + IteratorT, + ::cuda::std::int64_t, + SelectImpl::SelectPotentiallyInPlace>; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::select::dispatch( + storage, + bytes, + d_data, + static_cast(nullptr), + d_data, + d_num_selected_out, + NullType{}, + equality_op, + num_items, + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_keys_in`` and ``d_values_in`` with runs of key-value pairs with consecutive + //! equal-valued keys, only the first key and its value from each run is selectively copied + //! to ``d_keys_out`` and ``d_values_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The user-provided equality operator, `equality_op`, is used to determine whether keys are equivalent + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + *d_num_selected_out)`` + //! - ``[d_values_in, d_values_in + num_items)`` + //! - ``[d_values_out, d_values_out + *d_num_selected_out)`` + //! - ``[d_num_selected_out, d_num_selected_out + 1)`` + //! + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 8 + //! int *d_keys_in; // e.g., [0, 2, 2, 9, 5, 5, 5, 8] + //! int *d_values_in; // e.g., [1, 2, 3, 4, 5, 6, 7, 8] + //! int *d_keys_out; // e.g., [ , , , , , , , ] + //! int *d_values_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSelect::UniqueByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, + //! d_keys_out, d_values_out, d_num_selected_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DeviceSelect::UniqueByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, + //! d_keys_out, d_values_out, d_num_selected_out, num_items); + //! + //! // d_keys_out <-- [0, 2, 9, 5, 8] + //! // d_values_out <-- [1, 2, 4, 5, 8] + //! // d_num_selected_out <-- [5] + //! + //! @endrst + //! + //! @tparam KeyInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input keys @iterator + //! + //! @tparam ValueInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input values @iterator + //! + //! @tparam KeyOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected keys @iterator + //! + //! @tparam ValueOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected values @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EqualityOpT + //! **[inferred]** Type of equality_op + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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 + //! Pointer to the input sequence of keys + //! + //! @param[in] d_values_in + //! Pointer to the input sequence of values + //! + //! @param[out] d_keys_out + //! Pointer to the output sequence of selected keys + //! + //! @param[out] d_values_out + //! Pointer to the output sequence of selected values + //! + //! @param[out] d_num_selected_out + //! Pointer to the total number of items selected (i.e., length of `d_keys_out` or `d_values_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_keys_in` or `d_values_in`) + //! + //! @param[in] equality_op + //! Binary predicate to determine equality + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template < + typename KeyInputIteratorT, + typename ValueInputIteratorT, + typename KeyOutputIteratorT, + typename ValueOutputIteratorT, + typename NumSelectedIteratorT, + typename NumItemsT, + typename EqualityOpT, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0, + ::cuda::std::enable_if_t<::cuda::std::indirect_binary_predicate, + int> = 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t UniqueByKey( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_keys_in, + ValueInputIteratorT d_values_in, + KeyOutputIteratorT d_keys_out, + ValueOutputIteratorT d_values_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + EqualityOpT equality_op, + const EnvT& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE_IF(d_temp_storage, "cub::DeviceSelect::UniqueByKey"); + + using offset_t = detail::choose_offset_t; + + using default_policy_selector = + detail::unique_by_key::policy_selector_from_types, + detail::it_value_t>; + return detail::dispatch_with_env_and_tuning( + d_temp_storage, temp_storage_bytes, env, [&](auto policy_selector, void* storage, size_t& bytes, auto stream) { + return detail::unique_by_key::dispatch( + storage, + bytes, + d_keys_in, + d_values_in, + d_keys_out, + d_values_out, + d_num_selected_out, + equality_op, + static_cast(num_items), + stream, + policy_selector); + }); + } + + //! @rst + //! Given an input sequence ``d_keys_in`` and ``d_values_in`` with runs of key-value pairs with consecutive + //! equal-valued keys, only the first key and its value from each run is selectively copied + //! to ``d_keys_out`` and ``d_values_out``. + //! The total number of items selected is written to ``d_num_selected_out``. + //! + //! .. versionadded:: 2.2.0 + //! First appears in CUDA Toolkit 12.3. + //! + //! - The ``==`` equality operator is used to determine whether keys are equivalent + //! - Copies of the selected items are compacted into ``d_out`` and maintain + //! their original relative ordering. + //! - In-place operations are not supported. There must be no overlap between + //! any of the provided ranges: + //! + //! - ``[d_keys_in, d_keys_in + num_items)`` + //! - ``[d_keys_out, d_keys_out + *d_num_selected_out)`` + //! - ``[d_values_in, d_values_in + num_items)`` + //! - ``[d_values_out, d_values_out + *d_num_selected_out)`` + //! - ``[d_num_selected_out, d_num_selected_out + 1)`` + //! + //! - @devicestorage + //! + //! Snippet + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! The code snippet below illustrates the compaction of items selected from an ``int`` device vector. + //! + //! .. code-block:: c++ + //! + //! #include // or equivalently + //! + //! // Declare, allocate, and initialize device-accessible pointers + //! // for input and output + //! int num_items; // e.g., 8 + //! int *d_keys_in; // e.g., [0, 2, 2, 9, 5, 5, 5, 8] + //! int *d_values_in; // e.g., [1, 2, 3, 4, 5, 6, 7, 8] + //! int *d_keys_out; // e.g., [ , , , , , , , ] + //! int *d_values_out; // e.g., [ , , , , , , , ] + //! int *d_num_selected_out; // e.g., [ ] + //! ... + //! + //! // Determine temporary device storage requirements + //! void *d_temp_storage = nullptr; + //! size_t temp_storage_bytes = 0; + //! cub::DeviceSelect::UniqueByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, + //! d_keys_out, d_values_out, d_num_selected_out, num_items); + //! + //! // Allocate temporary storage + //! cudaMalloc(&d_temp_storage, temp_storage_bytes); + //! + //! // Run selection + //! cub::DeviceSelect::UniqueByKey( + //! d_temp_storage, temp_storage_bytes, + //! d_keys_in, d_values_in, + //! d_keys_out, d_values_out, d_num_selected_out, num_items); + //! + //! // d_keys_out <-- [0, 2, 9, 5, 8] + //! // d_values_out <-- [1, 2, 4, 5, 8] + //! // d_num_selected_out <-- [5] + //! + //! @endrst + //! + //! @tparam KeyInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input keys @iterator + //! + //! @tparam ValueInputIteratorT + //! **[inferred]** Random-access input iterator type for reading input values @iterator + //! + //! @tparam KeyOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected keys @iterator + //! + //! @tparam ValueOutputIteratorT + //! **[inferred]** Random-access output iterator type for writing selected values @iterator + //! + //! @tparam NumSelectedIteratorT + //! **[inferred]** Output iterator type for recording the number of items selected @iterator + //! + //! @tparam NumItemsT + //! **[inferred]** Type of num_items + //! + //! @tparam EnvT + //! **[inferred]** Environment type (e.g., `cuda::std::execution::env<...>`) + //! + //! @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 + //! Pointer to the input sequence of keys + //! + //! @param[in] d_values_in + //! Pointer to the input sequence of values + //! + //! @param[out] d_keys_out + //! Pointer to the output sequence of selected keys + //! + //! @param[out] d_values_out + //! Pointer to the output sequence of selected values + //! + //! @param[out] d_num_selected_out + //! Pointer to the total number of items selected (i.e., length of `d_keys_out` or `d_values_out`) + //! + //! @param[in] num_items + //! Total number of input items (i.e., length of `d_keys_in` or `d_values_in`) + //! + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template < + typename KeyInputIteratorT, + typename ValueInputIteratorT, + typename KeyOutputIteratorT, + typename ValueOutputIteratorT, + typename NumSelectedIteratorT, + typename NumItemsT, + typename EnvT = ::cuda::std::execution::env<>, + ::cuda::std::enable_if_t<::cuda::std::is_integral_v, int> = 0, + ::cuda::std::enable_if_t, int> = + 0> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t UniqueByKey( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_keys_in, + ValueInputIteratorT d_values_in, + KeyOutputIteratorT d_keys_out, + ValueOutputIteratorT d_values_out, + NumSelectedIteratorT d_num_selected_out, + NumItemsT num_items, + const EnvT& env = {}) + { + return UniqueByKey( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_values_in, + d_keys_out, + d_values_out, + d_num_selected_out, + num_items, + ::cuda::std::equal_to<>{}, + env); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_topk.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_topk.cuh new file mode 100644 index 00000000..d8c9b3ca --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_topk.cuh @@ -0,0 +1,2007 @@ +// 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 + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail +{ +template +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; + using out_offset_t = ::cuda::std:: + conditional_t), offset_t, choose_offset_t>; + + // Query environment properties to determine if the user-requested configuration is supported + static_assert(!::cuda::std::execution::__queryable_with, + "Determinism should be used inside requires to have an effect."); + using requirements_t = ::cuda::std::execution:: + __query_result_or_t>; + using requested_determinism_t = + ::cuda::std::execution::__query_result_or_t; + using requested_order_t = + ::cuda::std::execution::__query_result_or_t; + constexpr auto is_determinism_not_guaranteed = + ::cuda::std::is_same_v; + constexpr auto is_output_order_unsorted = + ::cuda::std::is_same_v; + + // 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>; + 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; + + return topk::dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + static_cast(num_items), + static_cast(k), + decomposer, + stream.get(), + policy_selector_t{}); +} + +template +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( + 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 `_, 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 `_. +//! +//! 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, 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( + 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, 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( + 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 > + CUB_RUNTIME_FUNCTION static // + ::cuda::std::enable_if_t, 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; + + static_assert(!detail::radix::can_twiddle, + "Custom decomposers are not supported for fundamental types; " + "use the non-decomposer API overload instead"); + + return detail::dispatch_topk( + 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, 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; + + static_assert(!detail::radix::can_twiddle, + "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( + 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, 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( + 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, 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( + 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 > + CUB_RUNTIME_FUNCTION static // + ::cuda::std::enable_if_t, 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; + + static_assert(!detail::radix::can_twiddle, + "Custom decomposers are not supported for fundamental types; " + "use the non-decomposer API overload instead"); + + return detail::dispatch_topk( + 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, 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; + + static_assert(!detail::radix::can_twiddle, + "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( + 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, 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( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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, 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( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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 > + CUB_RUNTIME_FUNCTION static // + ::cuda::std::enable_if_t, 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; + + static_assert(!detail::radix::can_twiddle, + "Custom decomposers are not supported for fundamental types; " + "use the non-decomposer API overload instead"); + + return detail::dispatch_topk( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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, 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; + + static_assert(!detail::radix::can_twiddle, + "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( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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, 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( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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, 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( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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 > + CUB_RUNTIME_FUNCTION static // + ::cuda::std::enable_if_t, 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; + + static_assert(!detail::radix::can_twiddle, + "Custom decomposers are not supported for fundamental types; " + "use the non-decomposer API overload instead"); + + return detail::dispatch_topk( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(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, 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; + + static_assert(!detail::radix::can_twiddle, + "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( + storage, + bytes, + d_keys_in, + d_keys_out, + static_cast(nullptr), + static_cast(nullptr), + num_items, + k, + decomposer, + env); + }); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/device_transform.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/device_transform.cuh new file mode 100644 index 00000000..cdbe48fa --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/device_transform.cuh @@ -0,0 +1,874 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#ifndef CCCL_DISABLE_NVRTC_COMPATIBILITY_CHECK +# if _CCCL_COMPILER(NVRTC) +# error \ + "Including is not supported when compiling with NVRTC. Include block-, warp-, or thread-level primitives instead (e.g. ). 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 +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail +{ +template +struct __return_constant +{ + T value; + + template + _CCCL_HOST_DEVICE auto operator()(Args&&...) const -> T + { + return value; + } +}; +} // namespace detail +CUB_NAMESPACE_END + +namespace cuda +{ +template +struct proclaims_copyable_arguments> : ::cuda::std::true_type +{}; +} // namespace cuda + +CUB_NAMESPACE_BEGIN +//! DeviceTransform provides device-wide, parallel operations for transforming elements tuple-wise from multiple input +//! sequences into an output sequence. +//! +//! @rst +//! +//! Tuning +//! +++++++++++++++++++++++++++++++++++++++++++++ +//! +//! All algorithms in DeviceTransform that accept an environment can be tuned by passing a custom :ref:`policy selector +//! ` that returns a :cpp:struct:`cub::TransformPolicy`, as shown in the example below: +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_transform_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin transform-policy-selector +//! :end-before: example-end transform-policy-selector +//! +//! .. literalinclude:: ../../../cub/test/catch2_test_device_transform_env_api.cu +//! :language: c++ +//! :dedent: +//! :start-after: example-begin transform-tuning +//! :end-before: example-end transform-tuning +//! @endrst +struct DeviceTransform +{ + template + CUB_RUNTIME_FUNCTION static cudaError_t __transform_internal( + ::cuda::std::tuple inputs, + RandomAccessIteratorOut output, + NumItemsT num_items, + Predicate predicate, + TransformOp transform_op, + const Env& env) + { + // We use int64_t internally, since it's faster than uint64_t and similar to a 32-bit offset type. See + // https://github.com/NVIDIA/cccl/issues/8805 for data. We use choose_signed_offset to just check if it can hold the + // value passed by the user, but otherwise ignore the chosen signed offset type. + using offset_t = ::cuda::std::int64_t; + if (const cudaError_t error = detail::choose_signed_offset::is_exceeding_offset_type(num_items)) + { + return error; + } + + const auto stream = ::cuda::__call_or(::cuda::get_stream, ::cuda::stream_ref{cudaStream_t{}}, env).get(); + + using tuning_env = + ::cuda::std::execution::__query_result_or_t>; + using default_policy_selector = + detail::transform::policy_selector_from_types, + ::cuda::std::tuple, + RandomAccessIteratorOut>; + + using policy_selector = + ::cuda::std::execution::__query_result_or_t; + +#if _CCCL_HAS_CONCEPTS() + static_assert(detail::transform::transform_policy_selector); +#endif // _CCCL_HAS_CONCEPTS() + + return detail::transform::dispatch( + ::cuda::std::move(inputs), + ::cuda::std::move(output), + static_cast(num_items), + ::cuda::std::move(predicate), + ::cuda::std::move(transform_op), + stream, + policy_selector{}); + } + + // TODO(bgruber): we want to eventually forward the output tuple to the kernel and optimize writing multiple streams + template + CUB_RUNTIME_FUNCTION static cudaError_t __transform_internal( + ::cuda::std::tuple inputs, + ::cuda::std::tuple outputs, + NumItemsT num_items, + Predicate predicate, + TransformOp transform_op, + const Env& env) + { + return __transform_internal( + ::cuda::std::move(inputs), + ::cuda::make_zip_iterator(::cuda::std::move(outputs)), + num_items, + ::cuda::std::move(predicate), + ::cuda::std::move(transform_op), + env); + } + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Transforms many input sequences into many output sequence, by applying a transformation operation on corresponding + //! input elements and writing the tuple result to the corresponding output elements. No guarantee is given on the + //! identity (i.e. address) of the objects passed to the call operator of the transformation operation. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_transform_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin transform-many-many + //! :end-before: example-end transform-many-many + //! + //! @endrst + //! + //! @param inputs A tuple of iterators to the input sequences where num_items elements are read from each. The + //! iterators' value types must be trivially relocatable. + //! @param outputs A tuple of iterators to the output sequences where num_items results are written to each. Each + //! sequence may point to the beginning of one of the input sequences, performing the transformation inplace. Any + //! output sequence must not overlap with any of the input sequence in any other way. + //! @param num_items The number of elements in each input and output sequence. + //! @param transform_op An n-ary function object, where n is the number of input sequences. The input iterators' value + //! types must be convertible to the parameters of the function object's call operator. The return type of the call + //! operator must be a tuple where each tuple element is assignable to the corresponding dereferenced output + //! iterators. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t Transform( + ::cuda::std::tuple inputs, + ::cuda::std::tuple outputs, + NumItemsT num_items, + TransformOp transform_op, + const Env& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceTransform::Transform"); + return __transform_internal( + ::cuda::std::move(inputs), + ::cuda::std::move(outputs), + num_items, + ::cuda::always_true{}, + ::cuda::std::move(transform_op), + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // Overload with additional parameters to specify temporary storage. Provided for compatibility with other CUB APIs. + template > + CUB_RUNTIME_FUNCTION static cudaError_t Transform( + void* d_temp_storage, + size_t& temp_storage_bytes, + ::cuda::std::tuple inputs, + ::cuda::std::tuple outputs, + NumItemsT num_items, + TransformOp transform_op, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return Transform( + ::cuda::std::move(inputs), ::cuda::std::move(outputs), num_items, ::cuda::std::move(transform_op), env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Transforms many input sequences into one output sequence, by applying a transformation operation on corresponding + //! input elements and writing the result to the corresponding output element. No guarantee is given on the identity + //! (i.e. address) of the objects passed to the call operator of the transformation operation. + //! + //! .. versionadded:: 2.8.0 + //! First appears in CUDA Toolkit 12.9. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_transform_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin transform-many + //! :end-before: example-end transform-many + //! + //! @endrst + //! + //! @param inputs A tuple of iterators to the input sequences where num_items elements are read from each. + //! @param output An iterator to the output sequence where num_items results are written to. May point to the + //! beginning of one of the input sequences, performing the transformation inplace. The output sequence must not + //! overlap with any of the input sequence in any other way. + //! @param num_items The number of elements in each input sequence. + //! @param transform_op An n-ary function object, where n is the number of input sequences. The input iterators' value + //! types must be convertible to the parameters of the function object's call operator. The return type of the call + //! operator must be assignable to the dereferenced output iterator. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t Transform( + ::cuda::std::tuple inputs, + RandomAccessIteratorOut output, + NumItemsT num_items, + TransformOp transform_op, + const Env& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceTransform::Transform"); + return __transform_internal( + ::cuda::std::move(inputs), + ::cuda::std::move(output), + num_items, + ::cuda::always_true{}, + ::cuda::std::move(transform_op), + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // Overload with additional parameters to specify temporary storage. Provided for compatibility with other CUB APIs. + template > + CUB_RUNTIME_FUNCTION static cudaError_t Transform( + void* d_temp_storage, + size_t& temp_storage_bytes, + ::cuda::std::tuple inputs, + RandomAccessIteratorOut output, + NumItemsT num_items, + TransformOp transform_op, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return Transform( + ::cuda::std::move(inputs), ::cuda::std::move(output), num_items, ::cuda::std::move(transform_op), env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Transforms one input sequence into one output sequence, by applying a transformation operation on each input + //! element and writing the result to the corresponding output element. No guarantee is given on the identity (i.e. + //! address) of the objects passed to the call operator of the transformation operation. + //! This is effectively calling Transform with a single-input tuple. + //! + //! .. versionadded:: 2.8.0 + //! First appears in CUDA Toolkit 12.9. + //! @endrst + //! + //! @param input An iterator to the input sequence where num_items elements are read from. + //! @param output An iterator to the output sequence where num_items results are written to. May point to the same + //! sequence as \p input, performing the transformation inplace. The output sequence must not overlap with the + //! input sequence in any other way. + //! @param num_items The number of elements in each input sequence. + //! @param transform_op A unary function object. The input iterator's value type must be convertible to the parameter + //! of the function object's call operator. The return type of the call operator must be assignable to the + //! dereferenced output iterator. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t Transform( + RandomAccessIteratorIn input, + RandomAccessIteratorOut output, + NumItemsT num_items, + TransformOp transform_op, + const Env& env = {}) + { + return Transform( + ::cuda::std::make_tuple(::cuda::std::move(input)), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(transform_op), + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // Overload with additional parameters to specify temporary storage. Provided for compatibility with other CUB APIs. + template > + CUB_RUNTIME_FUNCTION static cudaError_t Transform( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorIn input, + RandomAccessIteratorOut output, + NumItemsT num_items, + TransformOp transform_op, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return Transform( + ::cuda::std::make_tuple(::cuda::std::move(input)), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(transform_op), + env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Fills the output sequence by invoking a generator operation for each output element and writing the result to it. + //! This is effectively calling Transform with no input sequences. + //! + //! .. versionadded:: 2.8.0 + //! First appears in CUDA Toolkit 12.9. + //! @endrst + //! + //! @param output An iterator to the output sequence where num_items results are written to. + //! @param num_items The number of elements to write to the output sequence. + //! @param generator A nullary function object. The return type of the call operator must be assignable to the + //! dereferenced output iterator. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t + Generate(RandomAccessIteratorOut output, NumItemsT num_items, Generator generator, const Env& env = {}) + { + static_assert(::cuda::std::is_invocable_v, "The passed generator must be a nullary function object"); + static_assert( + ::cuda::std::is_assignable_v, + ::cuda::std::invoke_result_t>, + "The return value of the generator's call operator must be assignable to the dereferenced output iterator"); + + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceTransform::Generate"); + return __transform_internal( + ::cuda::std::make_tuple(), + ::cuda::std::move(output), + num_items, + ::cuda::always_true{}, + ::cuda::std::move(generator), + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // Overload with additional parameters to specify temporary storage. Provided for compatibility with other CUB APIs. + template > + CUB_RUNTIME_FUNCTION static cudaError_t Generate( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorOut output, + NumItemsT num_items, + Generator generator, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return Generate(::cuda::std::move(output), num_items, ::cuda::std::move(generator), env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Fills the output sequence by writing the provided value to each element of the output sequence. + //! This is effectively calling Generate with a functor returning that value. + //! + //! .. versionadded:: 2.8.0 + //! First appears in CUDA Toolkit 12.9. + //! @endrst + //! + //! @param output An iterator to the output sequence where num_items results are written to. + //! @param num_items The number of elements to write to the output sequence. + //! @param value The value to write. Must be assignable to the dereferenced output iterator. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t + Fill(RandomAccessIteratorOut output, NumItemsT num_items, Value value, const Env& env = {}) + { + static_assert(::cuda::std::is_assignable_v, Value>, + "The passed value must be assignable to the dereferenced output iterator"); + + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceTransform::Fill"); + return __transform_internal( + ::cuda::std::make_tuple(), + ::cuda::std::move(output), + num_items, + ::cuda::always_true{}, + detail::__return_constant{::cuda::std::move(value)}, + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // Overload with additional parameters to specify temporary storage. Provided for compatibility with other CUB APIs. + template > + CUB_RUNTIME_FUNCTION static cudaError_t + Fill(void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorOut output, + NumItemsT num_items, + Value value, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return Fill(::cuda::std::move(output), num_items, ::cuda::std::move(value), env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Selectively transforms many input sequences into one output sequence, by applying a transformation operation on + //! corresponding input elements, if a given predicate is true, and writing the result to the corresponding output + //! element. No guarantee is given on the identity (i.e. address) of the objects passed to the call operator of the + //! predicate and transformation operation. Output elements for which the predicate returns false are not written to. + //! + //! .. versionadded:: 2.8.0 + //! First appears in CUDA Toolkit 12.9. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_transform_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin transform-if + //! :end-before: example-end transform-if + //! + //! @endrst + //! + //! @param inputs A tuple of iterators to the input sequences where num_items elements are read from each. + //! @param output An iterator to the output sequence where num_items results are written to. May point to the + //! beginning of one of the input sequences, performing the transformation inplace. The output sequence must not + //! overlap with any of the input sequence in any other way. + //! @param num_items The number of elements in each input sequence. + //! @param predicate An n-ary function object, where n is the number of input sequences. The input iterators' value + //! types must be convertible to the parameters of the function object's call operator, which must return a boolean + //! value. + //! @param transform_op An n-ary function object, where n is the number of input sequences. The input iterators' value + //! types must be convertible to the parameters of the function object's call operator. The return type of the call + //! operator must be assignable to the dereferenced output iterator. Will only be invoked if \p predicate returns + //! true. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t TransformIf( + ::cuda::std::tuple inputs, + RandomAccessIteratorOut output, + NumItemsT num_items, + Predicate predicate, + TransformOp transform_op, + const Env& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceTransform::TransformIf"); + return __transform_internal( + ::cuda::std::move(inputs), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(predicate), + ::cuda::std::move(transform_op), + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // Overload with additional parameters to specify temporary storage. Provided for compatibility with other CUB APIs. + template > + CUB_RUNTIME_FUNCTION static cudaError_t TransformIf( + void* d_temp_storage, + size_t& temp_storage_bytes, + ::cuda::std::tuple inputs, + RandomAccessIteratorOut output, + NumItemsT num_items, + Predicate predicate, + TransformOp transform_op, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return TransformIf( + ::cuda::std::move(inputs), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(predicate), + ::cuda::std::move(transform_op), + env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Selectively transforms one input sequence into one output sequence, by applying a transformation operation on each + //! input element, if a given predicate is true, and writing the result to the corresponding output element. No + //! guarantee is given on the identity (i.e. address) of the objects passed to the call operator of the predicate and + //! transformation operation. Output elements for which the predicate returns false are not written to. + //! + //! .. versionadded:: 2.8.0 + //! First appears in CUDA Toolkit 12.9. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_transform_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin transform-if-single + //! :end-before: example-end transform-if-single + //! + //! @endrst + //! + //! @param input An iterator to the input sequence where num_items elements are read from. + //! @param output An iterator to the output sequence where num_items results are written to. May point to the same + //! sequence as \p input, performing the transformation inplace. The output sequence must not overlap with the + //! input sequence in any other way. + //! @param num_items The number of elements in each input sequence. + //! @param predicate A unary function objects returning \p bool. The input iterators' value types must be convertible + //! to the parameters of the function object's call operator. + //! @param transform_op A unary function object. The input iterator's value type must be convertible to the + //! parameter of the function object's call operator. The return type of the call operator must be assignable to the + //! dereferenced output iterator. Will only be invoked if \p predicate returns true. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t TransformIf( + RandomAccessIteratorIn input, + RandomAccessIteratorOut output, + NumItemsT num_items, + Predicate predicate, + TransformOp transform_op, + const Env& env = {}) + { + return TransformIf( + ::cuda::std::make_tuple(::cuda::std::move(input)), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(predicate), + ::cuda::std::move(transform_op), + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + // Overload with additional parameters to specify temporary storage. Provided for compatibility with other CUB APIs. + template > + CUB_RUNTIME_FUNCTION static cudaError_t TransformIf( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorIn input, + RandomAccessIteratorOut output, + NumItemsT num_items, + Predicate predicate, + TransformOp transform_op, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return TransformIf( + ::cuda::std::make_tuple(::cuda::std::move(input)), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(predicate), + ::cuda::std::move(transform_op), + env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Overview + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! Transforms many input sequences into one output sequence, by applying a transformation operation on corresponding + //! input elements and writing the result to the corresponding output element. The objects passed to the call operator + //! of the transformation operation are guaranteed to reside in the input sequences and are never copied. + //! + //! .. versionadded:: 2.8.0 + //! First appears in CUDA Toolkit 12.9. + //! + //! A Simple Example + //! +++++++++++++++++++++++++++++++++++++++++++++ + //! + //! .. literalinclude:: ../../../cub/test/catch2_test_device_transform_api.cu + //! :language: c++ + //! :dedent: + //! :start-after: example-begin transform-many-stable + //! :end-before: example-end transform-many-stable + //! + //! @endrst + //! + //! @param inputs A tuple of iterators to the input sequences where num_items elements are read from each. + //! @param output An iterator to the output sequence where num_items results are written to. May point to the + //! beginning of one of the input sequences, performing the transformation inplace. The output sequence must not + //! overlap with any of the input sequence in any other way. + //! @param num_items The number of elements in each input sequence. + //! @param transform_op An n-ary function object, where n is the number of input sequences. The input iterators' value + //! types must be convertible to the parameters of the function object's call operator. The return type of the call + //! operator must be assignable to the dereferenced output iterator. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t TransformStableArgumentAddresses( + ::cuda::std::tuple inputs, + RandomAccessIteratorOut output, + NumItemsT num_items, + TransformOp transform_op, + const Env& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceTransform::TransformStableArgumentAddresses"); + return __transform_internal( + ::cuda::std::move(inputs), + ::cuda::std::move(output), + num_items, + ::cuda::always_true{}, + ::cuda::std::move(transform_op), + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + template > + CUB_RUNTIME_FUNCTION static cudaError_t TransformStableArgumentAddresses( + void* d_temp_storage, + size_t& temp_storage_bytes, + ::cuda::std::tuple inputs, + RandomAccessIteratorOut output, + NumItemsT num_items, + TransformOp transform_op, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return TransformStableArgumentAddresses( + ::cuda::std::move(inputs), ::cuda::std::move(output), num_items, ::cuda::std::move(transform_op), env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + //! @rst + //! Transforms one input sequence into one output sequence, by applying a transformation operation on corresponding + //! input elements and writing the result to the corresponding output element. The objects passed to the call operator + //! of the transformation operation are guaranteed to reside in the input sequences and are never copied. + //! This is effectively calling TransformStableArgumentAddresses with a single-input tuple. + //! + //! .. versionadded:: 2.8.0 + //! First appears in CUDA Toolkit 12.9. + //! @endrst + //! + //! @param input An iterator to the input sequence where num_items elements are read from. + //! @param output An iterator to the output sequence where num_items results are written to. May point to the + //! beginning of one of the input sequences, performing the transformation inplace. The output sequence must not + //! overlap with any of the input sequence in any other way. + //! @param num_items The number of elements in each input sequence. + //! @param transform_op An n-ary function object, where n is the number of input sequences. The input iterators' value + //! types must be convertible to the parameters of the function object's call operator. The return type of the call + //! operator must be assignable to the dereferenced output iterator. + //! @param[in] env + //! **[optional]** Execution environment. Default is ``cuda::std::execution::env{}``. + template > + CUB_RUNTIME_FUNCTION static cudaError_t TransformStableArgumentAddresses( + RandomAccessIteratorIn input, + RandomAccessIteratorOut output, + NumItemsT num_items, + TransformOp transform_op, + const Env& env = {}) + { + return TransformStableArgumentAddresses( + ::cuda::std::make_tuple(::cuda::std::move(input)), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(transform_op), + env); + } + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + template > + CUB_RUNTIME_FUNCTION static cudaError_t TransformStableArgumentAddresses( + void* d_temp_storage, + size_t& temp_storage_bytes, + RandomAccessIteratorIn input, + RandomAccessIteratorOut output, + NumItemsT num_items, + TransformOp transform_op, + const EnvT& env = {}) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + return TransformStableArgumentAddresses( + ::cuda::std::make_tuple(::cuda::std::move(input)), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(transform_op), + env); + } +#endif // _CCCL_DOXYGEN_INVOKED + + // internal, used only by Thrust + template > + CUB_RUNTIME_FUNCTION static cudaError_t __transform_if_stable_argument_addresses( + ::cuda::std::tuple inputs, + RandomAccessIteratorOut output, + NumItemsT num_items, + Predicate predicate, + TransformOp transform_op, + const Env& env = {}) + { + _CCCL_NVTX_RANGE_SCOPE("cub::DeviceTransform::TransformIfStableArgumentAddresses"); + return __transform_internal( + ::cuda::std::move(inputs), + ::cuda::std::move(output), + num_items, + ::cuda::std::move(predicate), + ::cuda::std::move(transform_op), + env); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_adjacent_difference.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_adjacent_difference.cuh new file mode 100644 index 00000000..69480e69 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_adjacent_difference.cuh @@ -0,0 +1,467 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::adjacent_difference +{ +template +_CCCL_KERNEL_ATTRIBUTES void DeviceAdjacentDifferenceInitKernel( + const InputIteratorT first, InputT* const result, const OffsetT num_tiles, const int items_per_tile) +{ + const int tile_idx = static_cast(blockIdx.x * blockDim.x + threadIdx.x); + AgentDifferenceInitT::Process(tile_idx, first, result, num_tiles, items_per_tile); +} + +template +_CCCL_KERNEL_ATTRIBUTES void DeviceAdjacentDifferenceDifferenceKernel( + const InputIteratorT input, + InputT* const first_tile_previous, + const OutputIteratorT result, + DifferenceOpT difference_op, + const OffsetT num_items) +{ + static_assert(::cuda::std::is_empty_v); + static constexpr AdjacentDifferencePolicy policy = current_policy(); + using AdjacentDifferencePolicyT = + agent_adjacent_difference_policy; + + // It is OK to introspect the return type or parameter types of the + // `operator()` function of `__device__` extended lambda within device code. + using OutputT = ::cuda::std::invoke_result_t; + + using Agent = + AgentDifference; + + __shared__ typename Agent::TempStorage storage; + + Agent agent(storage, input, first_tile_previous, result, difference_op, num_items); + + int tile_idx = static_cast(blockIdx.x); + OffsetT tile_base = static_cast(tile_idx) * AdjacentDifferencePolicyT::ITEMS_PER_TILE; + + agent.Process(tile_idx, tile_base); +} + +template +struct policy_selector_from_hub +{ + // this is only called in device code, so we can ignore the cc parameter + _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> AdjacentDifferencePolicy + { + using p = typename PolicyHub::MaxPolicy::ActivePolicy::AdjacentDifferencePolicy; + return AdjacentDifferencePolicy{ + p::BLOCK_THREADS, p::ITEMS_PER_THREAD, p::LOAD_ALGORITHM, p::LOAD_MODIFIER, p::STORE_ALGORITHM}; + } +}; +} // namespace detail::adjacent_difference + +enum class ReadOption +{ + Left, + Right +}; + +// TODO(bgruber): remove in CCL 4.0 +//! Deprecated [Since 3.5] +template > +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceAdjacentDifference") DispatchAdjacentDifference +{ + using InputT = detail::it_value_t; + + void* d_temp_storage; + size_t& temp_storage_bytes; + InputIteratorT d_input; + OutputIteratorT d_output; + OffsetT num_items; + DifferenceOpT difference_op; + cudaStream_t stream; + + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchAdjacentDifference( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_input, + OutputIteratorT d_output, + OffsetT num_items, + DifferenceOpT difference_op, + cudaStream_t stream) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_input(d_input) + , d_output(d_output) + , num_items(num_items) + , difference_op(difference_op) + , stream(stream) + {} + + /// Invocation + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke() + { + using AdjacentDifferencePolicyT = typename ActivePolicyT::AdjacentDifferencePolicy; + + cudaError error = cudaSuccess; + + do + { + constexpr int tile_size = AdjacentDifferencePolicyT::ITEMS_PER_TILE; + const int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + size_t first_tile_previous_size = (AliasOpt == MayAlias::Yes) * num_tiles * sizeof(InputT); + + void* allocations[1] = {nullptr}; + size_t allocation_sizes[1] = {(AliasOpt == MayAlias::Yes) * first_tile_previous_size}; + + error = CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)); + + if (cudaSuccess != error) + { + break; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage + // allocation + + if (temp_storage_bytes == 0) + { + temp_storage_bytes = 1; + } + + break; + } + + if (num_items == OffsetT{}) + { + break; + } + + auto first_tile_previous = reinterpret_cast(allocations[0]); + + if constexpr (AliasOpt == MayAlias::Yes) + { + using AgentDifferenceInitT = + detail::adjacent_difference::AgentDifferenceInit; + + constexpr int init_block_size = AgentDifferenceInitT::BLOCK_THREADS; + const int init_grid_size = ::cuda::ceil_div(num_tiles, init_block_size); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceAdjacentDifferenceInitKernel" + "<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + init_block_size, + reinterpret_cast(stream)); +#endif // CUB_DEBUG_LOG + + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_block_size, 0, stream) + .doit(detail::adjacent_difference:: + DeviceAdjacentDifferenceInitKernel, + d_input, + first_tile_previous, + num_tiles, + tile_size)); + if (cudaSuccess != error) + { + break; + } + + error = CubDebug(detail::DebugSyncStream(stream)); + + if (cudaSuccess != error) + { + break; + } + } + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceAdjacentDifferenceDifferenceKernel" + "<<<%d, %d, 0, %lld>>>()\n", + num_tiles, + AdjacentDifferencePolicyT::BLOCK_THREADS, + reinterpret_cast(stream)); +#endif // CUB_DEBUG_LOG + + using KernelPolicySelector = detail::adjacent_difference::policy_selector_from_hub; + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + num_tiles, AdjacentDifferencePolicyT::BLOCK_THREADS, 0, stream) + .doit(detail::adjacent_difference::DeviceAdjacentDifferenceDifferenceKernel < KernelPolicySelector, + InputIteratorT, + OutputIteratorT, + DifferenceOpT, + OffsetT, + InputT, + AliasOpt == MayAlias::Yes, + ReadOpt == ReadOption::Left >, + d_input, + first_tile_previous, + d_output, + difference_op, + num_items)); + if (cudaSuccess != error) + { + break; + } + + error = CubDebug(detail::DebugSyncStream(stream)); + + if (cudaSuccess != error) + { + break; + } + + } while (false); + + return error; + } + + CUB_RUNTIME_FUNCTION static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_input, + OutputIteratorT d_output, + OffsetT num_items, + DifferenceOpT difference_op, + cudaStream_t stream) + { + cudaError error = cudaSuccess; + do + { + // Get PTX version + int ptx_version = 0; + error = CubDebug(PtxVersion(ptx_version)); + if (cudaSuccess != error) + { + break; + } + + // Create dispatch functor + DispatchAdjacentDifference dispatch( + d_temp_storage, temp_storage_bytes, d_input, d_output, num_items, difference_op, stream); + + // Dispatch to chained policy + error = CubDebug(PolicyHub::MaxPolicy::Invoke(ptx_version, dispatch)); + if (cudaSuccess != error) + { + break; + } + } while (false); + + return error; + } +}; + +namespace detail::adjacent_difference +{ +template , + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_input, + OutputIteratorT d_output, + NumItemsT num_items, + DifferenceOpT difference_op, + cudaStream_t stream, + TuningEnvT tuning_env = {}, + KernelLauncherFactory launcher_factory = {}) +{ + using input_t = detail::it_value_t; + using offset_t = detail::choose_offset_t; + using default_policy_selector_t = + detail::adjacent_difference::policy_selector_from_types; + using default_policy_t = decltype(default_policy_selector_t{}(::cuda::compute_capability{})); + + auto policy_selector = + ::cuda::std::execution::__query_or(tuning_env, default_policy_t{}, default_policy_selector_t{}); + using policy_selector_t = decltype(policy_selector); +#if _CCCL_HAS_CONCEPTS() + static_assert(adjacent_difference_policy_selector, + "Invalid policy_selector_t for adjacent_difference::dispatch"); +#endif // _CCCL_HAS_CONCEPTS() + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + const AdjacentDifferencePolicy active_policy = policy_selector(cc); +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceAdjacentDifference to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const int tile_size = active_policy.threads_per_block * active_policy.items_per_thread; + const int num_tiles = static_cast(::cuda::ceil_div(static_cast(num_items), tile_size)); + + size_t first_tile_previous_size = (AliasOpt == MayAlias::Yes) * num_tiles * sizeof(input_t); + + void* allocations[1] = {nullptr}; + size_t allocation_sizes[1] = {(AliasOpt == MayAlias::Yes) * first_tile_previous_size}; + + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + if (temp_storage_bytes == 0) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + if (static_cast(num_items) == offset_t{}) + { + return cudaSuccess; + } + + auto first_tile_previous = reinterpret_cast(allocations[0]); + + if constexpr (AliasOpt == MayAlias::Yes) + { + using AgentDifferenceInitT = AgentDifferenceInit; + + constexpr int init_block_size = AgentDifferenceInitT::BLOCK_THREADS; + const int init_grid_size = ::cuda::ceil_div(num_tiles, init_block_size); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceAdjacentDifferenceInitKernel" + "<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + init_block_size, + reinterpret_cast(stream)); +#endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_block_size, 0, stream) + .doit(detail::adjacent_difference:: + DeviceAdjacentDifferenceInitKernel, + d_input, + first_tile_previous, + num_tiles, + tile_size))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceAdjacentDifferenceDifferenceKernel" + "<<<%d, %d, 0, %lld>>>()\n", + num_tiles, + active_policy.threads_per_block, + reinterpret_cast(stream)); +#endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_tiles, active_policy.threads_per_block, 0, stream) + .doit(DeviceAdjacentDifferenceDifferenceKernel < policy_selector_t, + InputIteratorT, + OutputIteratorT, + DifferenceOpT, + offset_t, + input_t, + AliasOpt == MayAlias::Yes, + ReadOpt == ReadOption::Left >, + d_input, + first_tile_previous, + d_output, + difference_op, + static_cast(num_items)))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + return cudaSuccess; +} +} // namespace detail::adjacent_difference + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_batch_memcpy.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_batch_memcpy.cuh new file mode 100644 index 00000000..b86dd7a8 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_batch_memcpy.cuh @@ -0,0 +1,551 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! \file +//! cub::detail::batch_memcpy::dispatch provides device-wide, parallel operations for copying data from a number of +//! given source buffers to their corresponding destination buffer. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +enum class CopyAlg +{ + Memcpy, + Copy +}; + +namespace detail::batch_memcpy +{ +// Type used to specialize the kernel templates for indexing the buffers processed within a single kernel invocation +using per_invocation_buffer_offset_t = ::cuda::std::uint32_t; + +/** + * Initialization kernel for tile status initialization (multi-block) + */ +template +_CCCL_KERNEL_ATTRIBUTES void InitTileStateKernel( + BufferOffsetScanTileStateT buffer_offset_scan_tile_state, + BlockOffsetScanTileStateT block_offset_scan_tile_state, + const TileOffsetT num_tiles) +{ + // Initialize tile status + buffer_offset_scan_tile_state.InitializeStatus(num_tiles); + block_offset_scan_tile_state.InitializeStatus(num_tiles); +} + +/** + * Kernel that copies buffers that need to be copied by at least one (and potentially many) thread + * blocks. + */ +template +#if _CCCL_HAS_CONCEPTS() + requires batch_memcpy_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().lookback.large_buffer.threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void MultiBlockBatchMemcpyKernel( + const InputBufferIt input_buffer_it, + const OutputBufferIt output_buffer_it, + const BufferSizeIteratorT buffer_sizes, + const BufferTileOffsetItT buffer_tile_offsets, + TileT buffer_offset_tile, + const TileOffsetT last_tile_offset) +{ + static constexpr BatchedCopyLargeBufferPolicy policy = current_policy().lookback.large_buffer; + using BufferSizeT = it_value_t; + /// Internal load/store type. For byte-wise memcpy, a single-byte type + using AliasT = typename ::cuda::std::conditional_t, + lazy_trait>>::type; + /// Types of the input and output buffers + using InputBufferT = it_value_t; + using OutputBufferT = it_value_t; + + constexpr uint32_t BLOCK_THREADS = static_cast(policy.threads_per_block); + constexpr uint32_t ITEMS_PER_THREAD = static_cast(policy.bytes_per_thread); + constexpr BufferSizeT TILE_SIZE = BufferSizeT{BLOCK_THREADS} * ITEMS_PER_THREAD; + + BufferOffsetT num_blev_buffers = buffer_offset_tile.LoadValid(last_tile_offset); + + uint32_t tile_id = blockIdx.x; + + // No block-level buffers => we're done here + if (num_blev_buffers == 0) + { + return; + } + + // While there's still tiles of bytes from block-level buffers to copied + do + { + __shared__ BufferOffsetT block_buffer_id; + + // Make sure thread 0 does not overwrite the buffer id before other threads have finished with + // the prior iteration of the loop + __syncthreads(); + + // Binary search the buffer that this tile belongs to + if (threadIdx.x == 0) + { + block_buffer_id = UpperBound(buffer_tile_offsets, num_blev_buffers, tile_id) - 1; + } + + // Make sure thread 0 has written the buffer this thread block is assigned to + __syncthreads(); + + const BufferOffsetT buffer_id = block_buffer_id; + + // The relative offset of this tile within the buffer it's assigned to + BufferSizeT tile_offset_within_buffer = + static_cast(tile_id - buffer_tile_offsets[buffer_id]) * TILE_SIZE; + + // If the tile has already reached beyond the work of the end of the last buffer + if (buffer_id >= num_blev_buffers - 1 && tile_offset_within_buffer > buffer_sizes[buffer_id]) + { + return; + } + + // Tiny remainders are copied without vectorizing loads + if (buffer_sizes[buffer_id] - tile_offset_within_buffer <= 32) + { + BufferSizeT thread_offset = tile_offset_within_buffer + threadIdx.x; + for (int i = 0; i < ITEMS_PER_THREAD; i++) + { + if (thread_offset < buffer_sizes[buffer_id]) + { + const auto value = read_item < MemcpyOpt == CopyAlg::Memcpy, AliasT, + InputBufferT > (input_buffer_it[buffer_id], thread_offset); + write_item( + output_buffer_it[buffer_id], thread_offset, value); + } + thread_offset += BLOCK_THREADS; + } + } + else + { + copy_items( + input_buffer_it[buffer_id], + output_buffer_it[buffer_id], + (::cuda::std::min) (buffer_sizes[buffer_id] - tile_offset_within_buffer, TILE_SIZE), + tile_offset_within_buffer); + } + + tile_id += gridDim.x; + } while (true); +} + +/** + * @brief Kernel that copies data from a batch of given source buffers to their corresponding + * destination buffer. If a buffer's size is too large to be copied by a single thread block, that + * buffer is put into a queue of buffers that will get picked up later on, where multiple blocks + * collaborate on each of these buffers. All other buffers get copied straight away.o + * + * @param input_buffer_it [in] Iterator providing the pointers to the source memory buffers + * @param output_buffer_it [in] Iterator providing the pointers to the destination memory buffers + * @param buffer_sizes [in] Iterator providing the number of bytes to be copied for each pair of + * buffers + * @param num_buffers [in] The total number of buffer pairs + * @param blev_buffer_srcs [out] The source pointers of buffers that require block-level + * collaboration + * @param blev_buffer_dsts [out] The destination pointers of buffers that require block-level + * collaboration + * @param blev_buffer_sizes [out] The sizes of buffers that require block-level collaboration + * @param blev_buffer_scan_state [in,out] Tile states for the prefix sum over the count of buffers + * requiring block-level collaboration (to "stream compact" (aka "select") BLEV-buffers) + * @param blev_block_scan_state [in,out] Tile states for the prefix sum over the number of thread + * blocks getting assigned to each buffer that requires block-level collaboration + */ +template +#if _CCCL_HAS_CONCEPTS() + requires batch_memcpy_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().lookback.small_buffer.threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void BatchMemcpyKernel( + const InputBufferIt input_buffer_it, + const OutputBufferIt output_buffer_it, + const BufferSizeIteratorT buffer_sizes, + const BufferOffsetT num_buffers, + const BlevBufferSrcsOutItT blev_buffer_srcs, + const BlevBufferDstsOutItT blev_buffer_dsts, + const BlevBufferSizesOutItT blev_buffer_sizes, + const BlevBufferTileOffsetsOutItT blev_buffer_tile_offsets, + const BLevBufferOffsetTileState blev_buffer_scan_state, + const BLevBlockOffsetTileState blev_block_scan_state) +{ + static constexpr BatchedCopySmallBufferPolicy policy = current_policy().lookback.small_buffer; + + // TODO(bgruber): refactor this in C++20, when we can pass policy as NTTP + using agent_policy_t = agent_batch_memcpy_policy< + policy.threads_per_block, + policy.buffers_per_thread, + policy.bytes_per_thread, + policy.prefer_pow2_bits, + policy.block_level_tile_size, + policy.warp_level_threshold, + policy.block_level_threshold, + delay_constructor_t, + delay_constructor_t>; + + // Block-level specialization + using AgentBatchMemcpyT = AgentBatchMemcpy< + agent_policy_t, + InputBufferIt, + OutputBufferIt, + BufferSizeIteratorT, + BufferOffsetT, + BlevBufferSrcsOutItT, + BlevBufferDstsOutItT, + BlevBufferSizesOutItT, + BlevBufferTileOffsetsOutItT, + BlockOffsetT, + BLevBufferOffsetTileState, + BLevBlockOffsetTileState, + MemcpyOpt == CopyAlg::Memcpy>; + + // Shared memory for AgentBatchMemcpy + __shared__ typename AgentBatchMemcpyT::TempStorage temp_storage; + + // Process this block's tile of input&output buffer pairs + AgentBatchMemcpyT( + temp_storage, + input_buffer_it, + output_buffer_it, + buffer_sizes, + num_buffers, + blev_buffer_srcs, + blev_buffer_dsts, + blev_buffer_sizes, + blev_buffer_tile_offsets, + blev_buffer_scan_state, + blev_block_scan_state) + .ConsumeTile(blockIdx.x); +} + +//! @tparam BlockOffsetT Integer type large enough to hold any offset in [0, num_thread_blocks_launched) +//! @tparam InputBufferIt **[inferred]** Random-access input iterator type providing the pointers to the source memory +//! buffers +//! @tparam OutputBufferIt **[inferred]** Random-access input iterator type providing the pointers to the destination +//! memory buffers +//! @tparam BufferSizeIteratorT **[inferred]** Random-access input iterator type providing the number of bytes to be +//! copied for each pair of buffers +template +#if _CCCL_HAS_CONCEPTS() + requires batch_memcpy_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputBufferIt input_buffer_it, + OutputBufferIt output_buffer_it, + BufferSizeIteratorT buffer_sizes, + ::cuda::std::int64_t num_buffers, + cudaStream_t stream, + PolicySelectorT policy_selector = {}) +{ + using per_invocation_buffer_offset_t = detail::batch_memcpy::per_invocation_buffer_offset_t; + using BufferSizeT = cub::detail::it_value_t; + using BLevBufferOffsetTileState = cub::ScanTileState; + using BLevBlockOffsetTileState = cub::ScanTileState; + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(ptx_compute_cap(cc))) + { + return error; + } + const BatchedCopyPolicy active_policy = policy_selector(cc); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceBatchMemcpy to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + enum : uint32_t + { + // Memory for the source pointers of the buffers that require block-level collaboration + MEM_BLEV_BUFFER_SRCS = 0, + // Memory for the destination pointers of the buffers that require block-level collaboration + MEM_BLEV_BUFFER_DSTS, + // Memory for the block-level buffers' sizes + MEM_BLEV_BUFFER_SIZES, + // Memory to keep track of the assignment of thread blocks to block-level buffers + MEM_BLEV_BUFFER_TBLOCK, + // Memory for the tile states of the prefix sum over the number of buffers that require + // block-level collaboration + MEM_BLEV_BUFFER_SCAN_STATE, + // Memory for the scan tile states of the prefix sum over the number of thread block's + // assigned up to and including a certain block-level buffer + MEM_BLEV_BLOCK_SCAN_STATE, + MEM_NUM_ALLOCATIONS + }; + + constexpr BlockOffsetT init_kernel_threads = 128U; + const auto tile_size = static_cast(active_policy.lookback.small_buffer.threads_per_block) + * static_cast(active_policy.lookback.small_buffer.buffers_per_thread); + + constexpr auto max_num_buffers_per_invocation = ::cuda::std::int64_t{512 * 1024 * 1024}; + static_assert(max_num_buffers_per_invocation <= ::cuda::std::numeric_limits::max()); + const auto max_num_buffers = ::cuda::std::min(max_num_buffers_per_invocation, num_buffers); + const auto max_num_tiles = static_cast(::cuda::ceil_div(max_num_buffers, tile_size)); + + using BlevBufferSrcsOutT = + ::cuda::std::_If>; + using BlevBufferDstOutT = + ::cuda::std::_If>; + using BlevBufferSrcsOutItT = BlevBufferSrcsOutT*; + using BlevBufferDstsOutItT = BlevBufferDstOutT*; + using BlevBufferSizesOutItT = BufferSizeT*; + using BlevBufferTileOffsetsOutItT = BlockOffsetT*; + + temporary_storage::layout temporary_storage_layout; + + auto blev_buffer_srcs_slot = temporary_storage_layout.get_slot(MEM_BLEV_BUFFER_SRCS); + auto blev_buffer_dsts_slot = temporary_storage_layout.get_slot(MEM_BLEV_BUFFER_DSTS); + auto blev_buffer_sizes_slot = temporary_storage_layout.get_slot(MEM_BLEV_BUFFER_SIZES); + auto blev_buffer_block_slot = temporary_storage_layout.get_slot(MEM_BLEV_BUFFER_TBLOCK); + auto blev_buffer_scan_slot = temporary_storage_layout.get_slot(MEM_BLEV_BUFFER_SCAN_STATE); + auto blev_buffer_block_scan_slot = temporary_storage_layout.get_slot(MEM_BLEV_BLOCK_SCAN_STATE); + + auto blev_buffer_srcs_alloc = blev_buffer_srcs_slot->template create_alias(); + auto blev_buffer_dsts_alloc = blev_buffer_dsts_slot->template create_alias(); + auto blev_buffer_sizes_alloc = blev_buffer_sizes_slot->template create_alias(); + auto blev_buffer_block_alloc = blev_buffer_block_slot->template create_alias(); + auto blev_buffer_scan_alloc = blev_buffer_scan_slot->template create_alias(); + auto blev_block_scan_alloc = blev_buffer_block_scan_slot->template create_alias(); + + size_t buffer_offset_scan_storage = 0; + size_t blev_block_scan_storage = 0; + if (const auto error = CubDebug( + BLevBufferOffsetTileState::AllocationSize(static_cast(max_num_tiles), buffer_offset_scan_storage))) + { + return error; + } + if (const auto error = CubDebug( + BLevBlockOffsetTileState::AllocationSize(static_cast(max_num_tiles), blev_block_scan_storage))) + { + return error; + } + + blev_buffer_srcs_alloc.grow(max_num_buffers); + blev_buffer_dsts_alloc.grow(max_num_buffers); + blev_buffer_sizes_alloc.grow(max_num_buffers); + blev_buffer_block_alloc.grow(max_num_buffers); + blev_buffer_scan_alloc.grow(buffer_offset_scan_storage); + blev_block_scan_alloc.grow(blev_block_scan_storage); + + if (d_temp_storage == nullptr) + { + temp_storage_bytes = temporary_storage_layout.get_size(); + return cudaSuccess; + } + if (num_buffers == 0) + { + return cudaSuccess; + } + + if (const auto error = CubDebug(temporary_storage_layout.map_to_buffer(d_temp_storage, temp_storage_bytes))) + { + return error; + } + + BlevBufferSrcsOutItT d_blev_src_buffers = blev_buffer_srcs_alloc.get(); + BlevBufferDstsOutItT d_blev_dst_buffers = blev_buffer_dsts_alloc.get(); + BlevBufferSizesOutItT d_blev_buffer_sizes = blev_buffer_sizes_alloc.get(); + BlevBufferTileOffsetsOutItT d_blev_block_offsets = blev_buffer_block_alloc.get(); + + auto init_scan_states_kernel = + detail::batch_memcpy::InitTileStateKernel; + auto batch_memcpy_non_blev_kernel = detail::batch_memcpy::BatchMemcpyKernel< + PolicySelectorT, + InputBufferIt, + OutputBufferIt, + BufferSizeIteratorT, + per_invocation_buffer_offset_t, + BlevBufferSrcsOutItT, + BlevBufferDstsOutItT, + BlevBufferSizesOutItT, + BlevBufferTileOffsetsOutItT, + BlockOffsetT, + BLevBufferOffsetTileState, + BLevBlockOffsetTileState, + MemcpyOpt>; + auto multi_block_memcpy_kernel = detail::batch_memcpy::MultiBlockBatchMemcpyKernel< + PolicySelectorT, + per_invocation_buffer_offset_t, + BlevBufferSrcsOutItT, + BlevBufferDstsOutItT, + BlevBufferSizesOutItT, + BlevBufferTileOffsetsOutItT, + BLevBufferOffsetTileState, + BlockOffsetT, + MemcpyOpt>; + + const auto blev_threads_per_block = static_cast(active_policy.lookback.large_buffer.threads_per_block); + + int device_ordinal; + if (const auto error = CubDebug(cudaGetDevice(&device_ordinal))) + { + return error; + } + int sm_count; + if (const auto error = CubDebug(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_ordinal))) + { + return error; + } + + int batch_memcpy_blev_occupancy; + if (const auto error = + CubDebug(MaxSmOccupancy(batch_memcpy_blev_occupancy, multi_block_memcpy_kernel, blev_threads_per_block))) + { + return error; + } + const int batch_memcpy_blev_grid_size = + static_cast(sm_count * batch_memcpy_blev_occupancy * subscription_factor); + + const ::cuda::std::int64_t num_invocations = ::cuda::ceil_div(num_buffers, max_num_buffers_per_invocation); + + for (::cuda::std::int64_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const auto current_buffer_offset = invocation_index * max_num_buffers_per_invocation; + const auto num_current_buffers = + ::cuda::std::min(max_num_buffers_per_invocation, num_buffers - current_buffer_offset); + const auto num_current_tiles = static_cast(::cuda::ceil_div(num_current_buffers, tile_size)); + const auto init_grid_size = static_cast(::cuda::ceil_div(num_current_tiles, init_kernel_threads)); + const auto batch_memcpy_grid_size = num_current_tiles; + + BLevBufferOffsetTileState buffer_scan_tile_state; + if (const auto error = CubDebug(buffer_scan_tile_state.Init( + static_cast(num_current_tiles), blev_buffer_scan_alloc.get(), buffer_offset_scan_storage))) + { + return error; + } + + BLevBlockOffsetTileState block_scan_tile_state; + if (const auto error = CubDebug(block_scan_tile_state.Init( + static_cast(num_current_tiles), blev_block_scan_alloc.get(), blev_block_scan_storage))) + { + return error; + } + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_kernel_threads, 0, stream) + .doit(init_scan_states_kernel, buffer_scan_tile_state, block_scan_tile_state, num_current_tiles))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + batch_memcpy_grid_size, active_policy.lookback.small_buffer.threads_per_block, 0, stream) + .doit(batch_memcpy_non_blev_kernel, + input_buffer_it + current_buffer_offset, + output_buffer_it + current_buffer_offset, + buffer_sizes + current_buffer_offset, + static_cast(num_current_buffers), + d_blev_src_buffers, + d_blev_dst_buffers, + d_blev_buffer_sizes, + d_blev_block_offsets, + buffer_scan_tile_state, + block_scan_tile_state))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + batch_memcpy_blev_grid_size, blev_threads_per_block, 0, stream) + .doit(multi_block_memcpy_kernel, + d_blev_src_buffers, + d_blev_dst_buffers, + d_blev_buffer_sizes, + d_blev_block_offsets, + buffer_scan_tile_state, + batch_memcpy_grid_size - 1))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} +} // namespace detail::batch_memcpy + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_batched_topk.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_batched_topk.cuh new file mode 100644 index 00000000..f95538f0 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_batched_topk.cuh @@ -0,0 +1,370 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, 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 unordered data items residing within device-accessible memory. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::batched_topk +{ +// ----------------------------------------------------------------------------- +// Internal: wrap the compile-time select direction into a discrete param for dispatch +// ----------------------------------------------------------------------------- + +// The selection direction is compile-time only: callers pass `::cuda::args::constant`, which maps to a +// value-less static_discrete_param. Because the direction is fixed at compile time and carries no runtime value, it +// can never disagree with its only supported option, so dispatch can never silently degrade to a no-op. +template +[[nodiscard]] _CCCL_HOST_DEVICE auto wrap_select_direction(::cuda::args::constant) +{ + return params::static_discrete_param{}; +} + +// The selection direction is intentionally a compile-time constant: only `::cuda::args::constant` is +// accepted (the overload above maps it to a value-less static_discrete_param). This catch-all documents that +// deliberate limitation and rejects anything else (e.g. a runtime `detail::topk::select` or a per-segment iterator of +// directions) with a clear diagnostic. It is an intent/documentation guard rather than a user-facing one: callers +// reach the algorithm through the min/max device entry points (DeviceBatchedTopK::{Max,Min}{Keys,Pairs}), which +// construct the matching `constant` internally, so `dispatch` is only ever invoked with a direction we create. +template +[[nodiscard]] _CCCL_HOST_DEVICE auto wrap_select_direction(SelectDirectionT) +{ + static_assert(::cuda::std::__always_false_v, + "DeviceBatchedTopK currently supports only compile-time selection directions: the min/max entry " + "points (DeviceBatchedTopK::{Max,Min}{Keys,Pairs}) dispatch with a " + "::cuda::args::constant; runtime or per-segment directions are " + "intentionally not supported"); + // Unreachable (the static_assert above always fires); keeps the return type well-formed so the only diagnostic is + // the message above. + return params::static_discrete_param{}; +} + +// ----------------------------------------------------------------------------- +// Helper: turn a segment ID into the number of large-segment-agent tiles needed +// to cover that segment. Wrapped in a transform_iterator, this produces the +// per-segment tile counts that we exclusive-scan to obtain per-segment tile +// offsets. +// ----------------------------------------------------------------------------- +template +struct segment_size_to_tile_count_op +{ + SegmentSizeParameterT segment_sizes; + int large_segment_agent_tile_size; + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr TotalNumItemsValueType operator()(SegmentIndexT segment_id) const + { + return static_cast( + ::cuda::ceil_div(params::get_param(segment_sizes, segment_id), large_segment_agent_tile_size)); + } +}; + +// ----------------------------------------------------------------------------- +// Segmented Top-K Dispatch +// ----------------------------------------------------------------------------- + +//! @param 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 temp_storage_bytes Reference to size in bytes of `d_temp_storage` allocation +//! @param d_key_segments_it d_key_segments_it[segment_index] -> iterator to the input sequence of key data for segment +//! `segment_index` +//! @param d_key_segments_out_it d_key_segments_out_it[segment_index] -> iterator to the output sequence of key data for +//! segment `segment_index` +//! @param d_value_segments_it d_value_segments_it[segment_index] -> iterator to the input sequence of associated value +//! items for segment `segment_index`. When cub::NullType**, only keys are provided. +//! @param d_value_segments_out_it d_value_segments_out_it[segment_index] -> iterator to the output sequence of +//! associated value items for segment `segment_index` +//! @param segment_sizes Parameter providing segment sizes for each segment +//! @param k Parameter providing K for each segment +//! @param select_directions Parameter providing the selection direction for each segment +//! @param num_segments Number of segments +//! @param total_num_items_guarantee Allows the user to provide a guarantee on the upper bound of the total number of +//! items +template >, + it_value_t>, + ::cuda::std::int64_t, + ::cuda::args::__traits::highest>> +#if _CCCL_HAS_CONCEPTS() + requires batched_topk_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + SelectDirectionT select_direction, + NumSegmentsParameterT num_segments, + [[maybe_unused]] TotalNumItemsGuaranteeT total_num_items_guarantee, + cudaStream_t stream = nullptr, + [[maybe_unused]] PolicySelector policy_selector = {}) +{ + using large_segment_tile_offset_t = typename ::cuda::args::__traits::element_type; + + // Wrap the raw enum into the internal discrete param type + auto select_directions = wrap_select_direction(select_direction); + using SelectDirectionParameterT = decltype(select_directions); + + // Helper that determines (a) whether there's any one-worker-per-segment policy supporting the range of segment + // sizes and k, and (b) if so, which set of one-worker-per-segment policies to use + constexpr auto policy = find_smallest_covering_policy< + PolicySelector, + SegmentSizeParameterT, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + large_segment_tile_offset_t>::policy; + constexpr worker_policy worker_per_segment_policy = policy.worker_per_segment_policy; + constexpr multi_worker_policy multi_worker_per_segment_policy = policy.multi_worker_per_segment_policy; + + static constexpr int worker_per_segment_tile_size = + worker_per_segment_policy.threads_per_block * worker_per_segment_policy.items_per_thread; + static constexpr bool any_small_segments = + ::cuda::args::__traits::lowest <= worker_per_segment_tile_size; + static constexpr bool only_small_segments = + ::cuda::args::__traits::highest <= worker_per_segment_tile_size; + + // Allocation layout: + // only_small_segments: [0] dummy. + // any_small_segments && !only_small_segments (mixed): [0] tile offsets, [1] counters struct, + // [2] large-segment ids. + // !any_small_segments (large-only): [0] tile offsets, [1] segment-size transform-scan temp storage. + static constexpr int allocations_array_size = only_small_segments ? 1 : (any_small_segments ? 3 : 2); + size_t allocation_sizes[allocations_array_size] = {1}; + + using num_segments_val_t = typename ::cuda::args::__traits::element_type; + using counters_t = batched_topk_counters; + using segment_size_scan_offset_t = detail::choose_offset_t; + using segment_size_scan_input_op_t = + segment_size_to_tile_count_op; + static constexpr auto multi_worker_per_segment_tile_size = + multi_worker_per_segment_policy.threads_per_block * multi_worker_per_segment_policy.items_per_thread; + const segment_size_scan_input_op_t segment_size_scan_input_op{segment_sizes, multi_worker_per_segment_tile_size}; + // Transform iterator over [0, num_segments) producing the tile-count for each segment. + [[maybe_unused]] const auto segment_size_scan_input_it = ::cuda::transform_iterator( + ::cuda::counting_iterator{num_segments_val_t{0}}, segment_size_scan_input_op); + + if constexpr (!only_small_segments) + { + const auto num_segments_val = params::get_param(num_segments, 0); + // Scan output + allocation_sizes[0] = num_segments_val * sizeof(large_segment_tile_offset_t); + if constexpr (any_small_segments) + { + allocation_sizes[1] = sizeof(counters_t); + // Large segment ids for indirectly accessing the large segment parameters + allocation_sizes[2] = num_segments_val * sizeof(num_segments_val_t); + } + else + { + // Query the temporary storage requirement of the segment-size transform-scan. + if (const auto error = CubDebug(detail::scan::dispatch( + nullptr, + allocation_sizes[1], + segment_size_scan_input_it, + static_cast(nullptr), + ::cuda::std::plus<>{}, + detail::InputValue(large_segment_tile_offset_t{0}), + static_cast(num_segments_val), + stream))) + { + return error; + } + } + } + + // Compute allocation pointers into the single storage blob (or compute the necessary size of the blob) + void* allocations[allocations_array_size] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + // TODO (elstehle): support number of segments provided by device-accessible iterator + // Only uniform number of segments are supported (i.e., we need to resolve the number of segments on the host) + static_assert(::cuda::args::__traits::is_single_value, + "Only uniform segment sizes are currently supported."); + + if constexpr (any_small_segments) + { + if constexpr (!only_small_segments) + { + // Zero-initialize the counters struct that holds the large-segment queue length and the block retirement + // counter; both are read by the agent's atomic operations and must start at 0. + if (const auto error = CubDebug(cudaMemsetAsync(allocations[1], 0, sizeof(counters_t), stream))) + { + return error; + } + } + const int grid_dim = static_cast(params::get_param(num_segments, 0)); + constexpr int block_dim = worker_per_segment_policy.threads_per_block; + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(grid_dim, block_dim, 0, stream) + .doit( + device_segmented_topk_kernel< + PolicySelector, + KeyInputItItT, + KeyOutputItItT, + ValueInputItItT, + ValueOutputItItT, + SegmentSizeParameterT, + KParameterT, + SelectDirectionParameterT, + NumSegmentsParameterT, + large_segment_tile_offset_t>, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + only_small_segments ? nullptr : static_cast(allocations[1]), + only_small_segments ? nullptr : static_cast(allocations[2]), + only_small_segments ? nullptr : static_cast(allocations[0])))) + { + return error; + } + } + else + { + // No small segments: the small-kernel epilogue (which would otherwise produce the per-segment tile offsets) does + // not run. Compute the per-segment tile offsets directly via a transform-scan over all segment sizes. + // The large segment agent will either consume these offsets directly (segment_id -> tile offset) or, when going + // through the large-segment queue, via a transform iterator over `d_large_segments_ids` (level of indirection). + if (const auto error = CubDebug(detail::scan::dispatch( + allocations[1], + allocation_sizes[1], + segment_size_scan_input_it, + static_cast(allocations[0]), + ::cuda::std::plus<>{}, + detail::InputValue(large_segment_tile_offset_t{0}), + static_cast(params::get_param(num_segments, 0)), + stream))) + { + return error; + } + } + + if constexpr (!only_small_segments) + { + // TODO (elstehle): support larger number of segments through multiple kernel launches + // Depending on any_small_segments, we need to either: + // - Indirectly get the large segment parameters via the queued large segment IDs + // - Directly take the segment parameters since all segments are large + } + return CubDebug(detail::DebugSyncStream(stream)); +} +// Env-based dispatch function handling memory allocation as well. This is usually done by the device-layer, but there +// is no public API for segmented topk yet. +template > +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_with_env( + KeyInputItItT d_key_segments_it, + KeyOutputItItT d_key_segments_out_it, + ValueInputItItT d_value_segments_it, + ValueOutputItItT d_value_segments_out_it, + SegmentSizeParameterT segment_sizes, + KParameterT k, + SelectDirectionParameterT select_directions, + NumSegmentsParameterT num_segments, + TotalNumItemsGuaranteeT total_num_items_guarantee, + const EnvT& env = {}) +{ + using default_policy_selector = + policy_selector_from_types>, + it_value_t>, + ::cuda::std::int64_t, + ::cuda::args::__traits::highest>; + return detail::dispatch_with_env_and_tuning( + env, [&](auto policy_selector, void* d_temp_storage, size_t& temp_storage_bytes, cudaStream_t stream) { + return dispatch( + d_temp_storage, + temp_storage_bytes, + d_key_segments_it, + d_key_segments_out_it, + d_value_segments_it, + d_value_segments_out_it, + segment_sizes, + k, + select_directions, + num_segments, + total_num_items_guarantee, + stream, + policy_selector); + }); +} +} // namespace detail::batched_topk + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_common.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_common.cuh new file mode 100644 index 00000000..69396fcc --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_common.cuh @@ -0,0 +1,68 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +CUB_NAMESPACE_BEGIN + +//! @brief Options for specifying memory aliasing +enum class MayAlias +{ + Yes, + No +}; + +//! @brief Options for specifying sorting order. +enum class SortOrder +{ + Ascending, + Descending +}; + +//! @brief Options for specifying the behavior of the stream compaction algorithm. +enum class SelectImpl +{ + //! Stream compaction, discarding rejected items. It's required that memory of input and output are disjoint. + Select, + //! Stream compaction, discarding rejected items. Memory of the input may be identical to the memory of the output. + SelectPotentiallyInPlace, + //! Partition, keeping rejected items. It's required that memory of input and output are disjoint. + Partition +}; + +//! @brief Options for forcing inclusive prefix-scan even when initial value has been provided +enum class ForceInclusive +{ + Yes, + No +}; + +// Options for specifying selection direction (e.g., for top-k selection). +namespace detail::topk +{ +enum class select +{ + // Select the elements with the lowest values + min, + // Select the elements with the highest values + max +}; +} // namespace detail::topk + +namespace detail +{ +struct use_default +{}; +} // namespace detail + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_copy_mdspan.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_copy_mdspan.cuh new file mode 100644 index 00000000..a0f3df37 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_copy_mdspan.cuh @@ -0,0 +1,82 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +//! @file +#pragma once + +#include + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::copy_mdspan +{ +template +struct copy_mdspan_t +{ + MdspanIn mdspan_in; + MdspanOut mdspan_out; + + _CCCL_HOST_DEVICE_API copy_mdspan_t(MdspanIn mdspan_in, MdspanOut mdspan_out) + : mdspan_in{mdspan_in} + , mdspan_out{mdspan_out} + {} + + template + _CCCL_DEVICE_API _CCCL_FORCEINLINE void operator()(Idx, Indices... indices) + { + mdspan_out(indices...) = mdspan_in(indices...); + } +}; + +template > +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t +copy(::cuda::std::mdspan mdspan_in, + ::cuda::std::mdspan mdspan_out, + const EnvT& env = {}) +{ + if (mdspan_in.is_exhaustive() && mdspan_out.is_exhaustive() + && detail::have_same_strides(mdspan_in.mapping(), mdspan_out.mapping())) + { + return cub::DeviceTransform::__transform_internal( + ::cuda::std::make_tuple(mdspan_in.data_handle()), + mdspan_out.data_handle(), + mdspan_in.size(), + ::cuda::always_true{}, + ::cuda::std::identity{}, + env); + } + // TODO (fbusato): add ForEachInLayout when mdspan_in and mdspan_out have compatible layouts + // Compatible layouts could use more efficient iteration patterns + return cub::DeviceFor::__for_each_in_extents( + ::cuda::std::layout_right::mapping{mdspan_in.extents()}, copy_mdspan_t{mdspan_in, mdspan_out}, env); +} +} // namespace detail::copy_mdspan + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_find.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_find.cuh new file mode 100644 index 00000000..2a7a24de --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_find.cuh @@ -0,0 +1,221 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +//! @file +//! cub::DeviceFind provides device-wide, parallel operations for computing search across a sequence of data items +//! residing within device-accessible memory. + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::find +{ +template +__launch_bounds__(1) _CCCL_KERNEL_ATTRIBUTES void init_found_pos_pointer(ValueType* found_pos_ptr, OffsetT num_items) +{ + // we immediately trigger launching the find kernel, before waiting for a previous kernel + _CCCL_PDL_TRIGGER_NEXT_LAUNCH(); + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + *found_pos_ptr = num_items; +} + +template +#if _CCCL_HAS_CONCEPTS() + requires find_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().threads_per_block)) _CCCL_KERNEL_ATTRIBUTES void find_kernel( + IteratorT d_in, OffsetT num_items, OffsetT* found_pos_ptr, PredicateT predicate) +{ + constexpr FindIfPolicy policy = current_policy(); + using agent_find_t = + agent_t; + + __shared__ typename agent_find_t::TempStorage sresult; + + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + agent_find_t{sresult.Alias(), d_in, predicate, found_pos_ptr, num_items}.Process(); +} + +template +__launch_bounds__(1) + _CCCL_KERNEL_ATTRIBUTES void copy_final_result_to_output_iterator(ValueType* found_pos_ptr, OutputIteratorT d_out) +{ + _CCCL_PDL_GRID_DEPENDENCY_SYNC(); + *d_out = *found_pos_ptr; +} + +template >> +#if _CCCL_HAS_CONCEPTS() + requires find_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + PredicateT predicate, + cudaStream_t stream, + PolicySelector policy_selector = {}) +{ + using output_t = it_value_t; + + // if the output iterator can be turned into a pointer, the value type is integral, and has the same size as OffsetT + // (we tolerate a sign mismatch, because both the output value type and the offset type must be able to represent all + // offsets), then we can just atomically write to the output pointer directly. + static constexpr bool can_write_to_output_directly = + THRUST_NS_QUALIFIER::is_contiguous_iterator_v && ::cuda::std::is_integral_v + && size_of == sizeof(OffsetT); + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(ptx_compute_cap(cc))) + { + return error; + } + + const FindIfPolicy active_policy = policy_selector(cc); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceFind to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const int tile_size = active_policy.threads_per_block * active_policy.items_per_thread; + const int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + int device_ordinal; + if (const auto error = CubDebug(cudaGetDevice(&device_ordinal))) + { + return error; + } + + int sm_count; + if (const auto error = CubDebug(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_ordinal))) + { + return error; + } + + using unwrapped_input_iterator_t = THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t; + auto kernel_ptr = find_kernel; + + int find_if_sm_occupancy; + if (const auto error = + CubDebug(cub::MaxSmOccupancy(find_if_sm_occupancy, kernel_ptr, active_policy.threads_per_block))) + { + return error; + } + + // no * CUB_SUBSCRIPTION_FACTOR(0) because max_blocks gets too big + const int max_blocks = find_if_sm_occupancy * sm_count; + const int findif_grid_size = ::cuda::std::min(num_tiles, max_blocks); + + // Temporary storage allocation requirements + void* allocations[1] = {}; + size_t allocation_sizes[1] = {sizeof(OffsetT)}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { // Return if the caller is simply requesting the size of the storage allocation + return cudaSuccess; + } + + OffsetT* found_pos_ptr = [&] { + if constexpr (can_write_to_output_directly) + { + return reinterpret_cast(THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out)); + } + else + { + return static_cast(allocations[0]); + } + }(); + + // use d_temp_storage as the intermediate device result to read and write from. Then store the final result in the + // output iterator. + if (const auto error = CubDebug(THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + 1, 1, 0, stream, /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(init_found_pos_pointer, found_pos_ptr, num_items))) + { + return error; + } + + // Unwrap the input iterator to convert device_ptr to T* (raw pointer). This ensures that dereferencing yields T& + // instead of device_reference, which is necessary for predicates that don't accept proxy types. + auto d_in_unwrapped = THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in); + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + findif_grid_size, + active_policy.threads_per_block, + 0, + stream, + /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(kernel_ptr, d_in_unwrapped, num_items, found_pos_ptr, predicate))) + { + return error; + } + + if constexpr (!can_write_to_output_directly) + { + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + 1, 1, 0, stream, /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(copy_final_result_to_output_iterator, found_pos_ptr, d_out))) + { + return error; + } + } + + // Sync the stream if specified to flush runtime errors + return CubDebug(detail::DebugSyncStream(stream)); +} +} // namespace detail::find +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_find_bound_sorted_values.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_find_bound_sorted_values.cuh new file mode 100644 index 00000000..86361ee7 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_find_bound_sorted_values.cuh @@ -0,0 +1,219 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +//! @file +//! Dispatch for cub::DeviceFind::LowerBoundSortedValues / UpperBoundSortedValues (merge-path partition + per-tile +//! search). + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) +# include +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + +CUB_NAMESPACE_BEGIN + +namespace detail::find_bound_sorted_values +{ +// Computes the merge path intersections at equally wide intervals (Odeh et al, IPDPS 2012). +template +_CCCL_KERNEL_ATTRIBUTES void device_partition_find_bound_sorted_values_kernel( + _CCCL_GRID_CONSTANT const HaystackIt d_range, + _CCCL_GRID_CONSTANT const Offset range_count, + _CCCL_GRID_CONSTANT const NeedlesIt d_values, + _CCCL_GRID_CONSTANT const Offset values_count, + _CCCL_GRID_CONSTANT const Offset num_diagonals, + _CCCL_GRID_CONSTANT Offset* const range_beg_offsets, + PartitionCompOp partition_comp) +{ + constexpr FindBoundSortedValuesPolicy policy = current_policy(); + constexpr int tile_size = policy.threads_per_block * policy.items_per_thread; + + const Offset diagonal_idx = static_cast(blockDim.x) * blockIdx.x + threadIdx.x; + if (diagonal_idx < num_diagonals) + { + const Offset diagonal = ::cuda::std::min(diagonal_idx * static_cast(tile_size), range_count + values_count); + range_beg_offsets[diagonal_idx] = + cub::MergePath(d_range, d_values, range_count, values_count, diagonal, partition_comp); + } +} + +template +__launch_bounds__(int(current_policy().threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void device_find_bound_sorted_values_kernel( + _CCCL_GRID_CONSTANT const HaystackIt d_range, + _CCCL_GRID_CONSTANT const NeedlesIt d_values, + _CCCL_GRID_CONSTANT const OutputIt d_output, + _CCCL_GRID_CONSTANT const Offset range_count, + _CCCL_GRID_CONSTANT const Offset values_count, + _CCCL_GRID_CONSTANT Offset* const range_beg_offsets, + CompareOp comp) +{ + constexpr FindBoundSortedValuesPolicy policy = current_policy(); + using AgentT = + agent_t; + + __shared__ typename AgentT::TempStorage temp_storage; + + AgentT{temp_storage.Alias(), d_range, d_values, d_output, range_count, values_count, range_beg_offsets, comp}(); +} + +template , it_value_t>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires find_bound_sorted_values_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + HaystackIt d_range, + Offset range_count, + NeedlesIt d_values, + Offset values_count, + OutputIt d_output, + CompareOp comp, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelLauncherFactory launcher_factory = {}) +{ + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + const auto active_policy = policy_selector(cc); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET( + NV_IS_HOST, + (::std::stringstream ss; ss << active_policy; _CubLog( + "Dispatching find_bound_sorted_values (merge-path) to arch %d with tuning: %s\n", cc.get(), ss.str().c_str());)) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const Offset tile_size = static_cast(active_policy.threads_per_block) * active_policy.items_per_thread; + if (range_count > cuda::std::numeric_limits::max() - values_count) + { + return cudaErrorInvalidValue; + } + const Offset total_items = range_count + values_count; + const Offset num_tiles = ::cuda::ceil_div(total_items, tile_size); + const Offset num_diagonals = num_tiles + 1; + + void* allocations[1] = {nullptr}; + const size_t allocation_sizes[1] = {static_cast(num_diagonals) * sizeof(Offset)}; + if (const auto error = + CubDebug(cub::detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr || num_tiles == 0) + { + return cudaSuccess; + } + + auto* range_beg_offsets = static_cast(allocations[0]); + + auto partition_comp = Mode::make_partition_comp(comp); + + { + // Lightweight pass; not worth exposing through the tuning system. + constexpr int threads_per_partition_block = 256; + const int partition_grid_size = + static_cast(::cuda::ceil_div(num_diagonals, Offset{threads_per_partition_block})); + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + partition_grid_size, threads_per_partition_block, 0, stream) + .doit(device_partition_find_bound_sorted_values_kernel>, + d_range, + range_count, + d_values, + values_count, + num_diagonals, + range_beg_offsets, + partition_comp))) + { + return error; + } + if (const auto error = CubDebug(cub::detail::DebugSyncStream(stream))) + { + return error; + } + } + + { + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + static_cast(num_tiles), active_policy.threads_per_block, 0, stream) + .doit( + device_find_bound_sorted_values_kernel, + d_range, + d_values, + d_output, + range_count, + values_count, + range_beg_offsets, + comp))) + { + return error; + } + if (const auto error = CubDebug(cub::detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} +} // namespace detail::find_bound_sorted_values + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_for.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_for.cuh new file mode 100644 index 00000000..77e03346 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_for.cuh @@ -0,0 +1,154 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::for_each +{ +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t +invoke_dynamic_block_size(OffsetT num_items, OpT op, cudaStream_t stream, ForPolicy active_policy) +{ + int threads_per_block = 256; + auto kernel = detail::for_each::dynamic_kernel; + NV_IF_TARGET(NV_IS_HOST, ({ + int _{}; + if (const auto error = CubDebug(cudaOccupancyMaxPotentialBlockSize(&_, &threads_per_block, kernel))) + { + return error; + } + })); + + const auto tile_size = static_cast(threads_per_block * active_policy.items_per_thread); + const auto num_tiles = ::cuda::ceil_div(num_items, tile_size); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking detail::for_each::dynamic_kernel<<<%d, %d, 0, %lld>>>(), " + "%d items per thread\n", + static_cast(num_tiles), + static_cast(threads_per_block), + reinterpret_cast(stream), + static_cast(active_policy.items_per_thread)); +#endif + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + static_cast(num_tiles), static_cast(threads_per_block), 0, stream) + .doit(kernel, num_items, op))) + { + return error; + } + + if (auto error = CubDebug(detail::DebugSyncStream(stream))) + { + CubDebug(error = SyncStream(stream)); // TODO(bgruber): this does not make sense to me + return error; + } + + return cudaSuccess; +} + +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t +invoke_static_block_size(OffsetT num_items, OpT op, cudaStream_t stream, ForPolicy active_policy) +{ + const int threads_per_block = active_policy.threads_per_block; + const int items_per_thread = active_policy.items_per_thread; + const auto tile_size = static_cast(threads_per_block) * static_cast(items_per_thread); + const auto num_tiles = ::cuda::ceil_div(num_items, tile_size); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking detail::for_each::static_kernel<<<%d, %d, 0, %lld>>>(), " + "%d items per thread\n", + static_cast(num_tiles), + static_cast(threads_per_block), + reinterpret_cast(stream), + static_cast(items_per_thread)); +#endif + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + static_cast(num_tiles), static_cast(threads_per_block), 0, stream) + .doit(detail::for_each::static_kernel, num_items, op))) + { + return error; + } + + if (auto error = CubDebug(detail::DebugSyncStream(stream))) + { + CubDebug(error = SyncStream(stream)); // TODO(bgruber): this does not make sense to me + return error; + } + + return cudaSuccess; +} + +// The dispatch layer is in the detail namespace until we figure out tuning API +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t +dispatch(OffsetT num_items, OpT op, cudaStream_t stream, PolicySelector policy_selector = {}) +{ + if (num_items == 0) + { + return cudaSuccess; + } + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(ptx_compute_cap(cc))) + { + return error; + } + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << policy_selector(cc); + _CubLog("Dispatching DeviceFor to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + return CubDebug(dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) { + constexpr ForPolicy active_policy = policy_getter(); + if constexpr (active_policy.threads_per_block > 0) + { + return invoke_static_block_size(num_items, op, stream, active_policy); + } + else + { + return invoke_dynamic_block_size(num_items, op, stream, active_policy); + } + })); +} +} // namespace detail::for_each + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_histogram.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_histogram.cuh new file mode 100644 index 00000000..704c57ad --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_histogram.cuh @@ -0,0 +1,1400 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::DeviceHistogram provides device-wide parallel operations for constructing histogram(s) + * from a sequence of samples data residing within device-accessible memory. + */ + +#pragma once + +#include + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::histogram +{ +// Maximum number of bins per channel for which we will use a privatized smem strategy +static constexpr int max_privatized_smem_bins = 256; + +template +struct DeviceHistogramKernelSource +{ + using TransformsT = detail::histogram::Transforms; + + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramInitKernel() + { + return &DeviceHistogramInitKernel; + } + + /// Returns the default histogram sweep kernel that receives pre-initialized decode operators from the host. + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramSweepKernel() + { + return &DeviceHistogramSweepKernel< + PolicyT, + PRIVATIZED_SMEM_BINS, + NUM_CHANNELS, + NUM_ACTIVE_CHANNELS, + SampleIteratorT, + CounterT, + PrivatizedDecodeOpT, + OutputDecodeOpT, + OffsetT>; + } + + /// Returns the device-init histogram sweep kernel that initializes decode operators from level arrays in the kernel. + template + _CCCL_HIDE_FROM_ABI CUB_RUNTIME_FUNCTION static constexpr auto HistogramSweepKernelDeviceInit() + { + // For DispatchEven, we use the scale transform to convert samples to + // privatized bins and pass-thru transform to convert privatized bins to + // output bins, vice verse for byte samples. + + // For DispatchRange, we use the search transform to convert samples to + // privatized bins and scale transform to convert privatized bins to output bins, + // vice verse for byte samples. + + using DecodeOpT = ::cuda::std::conditional_t>; + + using PrivatizedDecodeOpT = + ::cuda::std::conditional_t; + using OutputDecodeOpT = + ::cuda::std::conditional_t; + + return &DeviceHistogramSweepDeviceInitKernel< + PolicyT, + PRIVATIZED_SMEM_BINS, + NUM_CHANNELS, + NUM_ACTIVE_CHANNELS, + SampleIteratorT, + CounterT, + FirstLevelArrayT, + SecondLevelArrayT, + PrivatizedDecodeOpT, + OutputDecodeOpT, + OffsetT, + IsEven>; + } + + CUB_RUNTIME_FUNCTION static constexpr size_t CounterSize() + { + return sizeof(CounterT); + } + + template + CUB_RUNTIME_FUNCTION static constexpr bool MayOverflow( + [[maybe_unused]] NumBinsT num_bins, + [[maybe_unused]] const UpperLevelArrayT& upper_level, + [[maybe_unused]] const LowerLevelArrayT& lower_level, + [[maybe_unused]] int channel) + { + using CommonT = typename TransformsT::ScaleTransform::CommonT; + + if constexpr (::cuda::std::is_integral_v) + { + using IntArithmeticT = typename TransformsT::ScaleTransform::IntArithmeticT; + return static_cast(upper_level[channel] - lower_level[channel]) + > (::cuda::std::numeric_limits::max() / static_cast(num_bins)); + } + else + { + return false; + } + } +}; + +template +#if _CCCL_HAS_CONCEPTS() + requires histogram_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_output_histograms, + ::cuda::std::array num_privatized_levels, + ::cuda::std::array num_output_levels, + FirstLevelArrayT first_level_array, + SecondLevelArrayT second_level_array, + int max_num_output_bins, + OffsetT num_row_pixels, + OffsetT num_rows, + OffsetT row_stride_samples, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + const HistogramPolicy active_policy = policy_selector(cc); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceHistogram to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const auto init_kernel = kernel_source.template HistogramInitKernel(); + auto sweep_kernel = [&] { + if constexpr (IsDeviceInit) + { + return kernel_source.template HistogramSweepKernelDeviceInit< + PolicySelector, + PRIVATIZED_SMEM_BINS, + FirstLevelArrayT, + SecondLevelArrayT, + IsEven, + IsByteSample>(); + } + else + { + using output_decode_op_t = typename FirstLevelArrayT::value_type; + using privatized_decode_op_t = typename SecondLevelArrayT::value_type; + return kernel_source + .template HistogramSweepKernel(); + } + }(); + + const int threads_per_block = active_policy.threads_per_block; + const int pixels_per_thread = active_policy.pixels_per_thread; + + // Get SM count + int sm_count; + if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) + { + return error; + } + + // Get SM occupancy for sweep_kernel + int histogram_sweep_sm_occupancy; + if (const auto error = + CubDebug(launcher_factory.MaxSmOccupancy(histogram_sweep_sm_occupancy, sweep_kernel, threads_per_block))) + { + return error; + } + + // Get device occupancy for sweep_kernel + int histogram_sweep_occupancy = histogram_sweep_sm_occupancy * sm_count; + + if (num_row_pixels * NUM_CHANNELS == row_stride_samples) + { + // Treat as a single linear array of samples + num_row_pixels *= num_rows; + num_rows = 1; + row_stride_samples = num_row_pixels * NUM_CHANNELS; + } + + // Get grid dimensions, trying to keep total blocks ~histogram_sweep_occupancy + int pixels_per_tile = threads_per_block * pixels_per_thread; + int tiles_per_row = static_cast(::cuda::ceil_div(num_row_pixels, pixels_per_tile)); + int blocks_per_row = ::cuda::std::min(histogram_sweep_occupancy, tiles_per_row); + int blocks_per_col = + (blocks_per_row > 0) + ? int(::cuda::std::min(static_cast(histogram_sweep_occupancy / blocks_per_row), num_rows)) + : 0; + int num_thread_blocks = blocks_per_row * blocks_per_col; + + dim3 sweep_grid_dims; + sweep_grid_dims.x = (unsigned int) blocks_per_row; + sweep_grid_dims.y = (unsigned int) blocks_per_col; + sweep_grid_dims.z = 1; + + // Temporary storage allocation requirements + constexpr int NUM_ALLOCATIONS = NUM_ACTIVE_CHANNELS + 1; + void* allocations[NUM_ALLOCATIONS] = {}; + size_t allocation_sizes[NUM_ALLOCATIONS]; + + for (int CHANNEL = 0; CHANNEL < NUM_ACTIVE_CHANNELS; ++CHANNEL) + { + allocation_sizes[CHANNEL] = + size_t(num_thread_blocks) * (num_privatized_levels[CHANNEL] - 1) * kernel_source.CounterSize(); + } + + allocation_sizes[NUM_ALLOCATIONS - 1] = GridQueue::AllocationSize(); + + // Alias the temporary allocations from the single storage blob (or compute the + // necessary size of the blob) + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage allocation + return cudaSuccess; + } + + // Construct the grid queue descriptor + GridQueue tile_queue(allocations[NUM_ALLOCATIONS - 1]); + + // Wrap arrays so we can pass them by-value to the kernel + ::cuda::std::array d_privatized_histograms_wrapper; + ::cuda::std::array num_privatized_bins_wrapper; + ::cuda::std::array num_output_bins_wrapper; + + auto* typed_allocations = reinterpret_cast(allocations); + ::cuda::std::copy(typed_allocations, typed_allocations + NUM_ACTIVE_CHANNELS, d_privatized_histograms_wrapper.begin()); + + auto minus_one = ::cuda::proclaim_return_type([](int levels) { + return levels - 1; + }); + ::cuda::std::transform( + num_privatized_levels.begin(), num_privatized_levels.end(), num_privatized_bins_wrapper.begin(), minus_one); + ::cuda::std::transform(num_output_levels.begin(), num_output_levels.end(), num_output_bins_wrapper.begin(), minus_one); + + constexpr int histogram_init_threads_per_block = 256; + int histogram_init_grid_dims = + (max_num_output_bins + histogram_init_threads_per_block - 1) / histogram_init_threads_per_block; + +// Log DeviceHistogramInitKernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceHistogramInitKernel<<<%d, %d, 0, %lld>>>()\n", + histogram_init_grid_dims, + histogram_init_threads_per_block, + (long long) stream); +#endif // CUB_DEBUG_LOG + + // Invoke histogram_init_kernel + if (const auto error = CubDebug( + launcher_factory(histogram_init_grid_dims, + histogram_init_threads_per_block, + 0, + stream, + /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(init_kernel, num_output_bins_wrapper, d_output_histograms, tile_queue))) + { + return error; + } + + // Return if empty problem + if (blocks_per_row == 0 || blocks_per_col == 0) + { + return cudaSuccess; + } + +// Log histogram_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking histogram_sweep_kernel<<<{%d, %d, %d}, %d, 0, %lld>>>(), %d pixels " + "per thread, %d SM occupancy\n", + sweep_grid_dims.x, + sweep_grid_dims.y, + sweep_grid_dims.z, + threads_per_block, + (long long) stream, + pixels_per_thread, + histogram_sweep_sm_occupancy); +#endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + launcher_factory( + sweep_grid_dims, threads_per_block, 0, stream, /* dependent_launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(sweep_kernel, + d_samples, + num_output_bins_wrapper, + num_privatized_bins_wrapper, + d_output_histograms, + d_privatized_histograms_wrapper, + first_level_array, + second_level_array, + num_row_pixels, + num_rows, + row_stride_samples, + tiles_per_row, + tile_queue))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + return cudaSuccess; +} + +// Dispatch routines for device-side decode operator initialization. These differ from the default dispatch routines in +// that they initialize the decode operators inside the kernel from level arrays, instead of initializing them on the +// host, but they are otherwise the same. This is needed for c.parallel, since we cannot instantiate the Transforms +// class on the host, as SampleT and LevelT are type erased. Another change needed is that the level arrays are now +// templates instead of concrete ::cuda::std::array types, since we are passing indirect_args from c.parallel. +// +// Initializing the decode operators inside the kernel results in some regressions (and some performance improvements) +// in the benchmark, which indicates that we need to re-tune the algorithm. This is why we kept the two dispatch paths +// (host init and device init) separate. We should think about merging them back together later on. + +/** + * Dispatch routine for HistogramEven with device-side decode operator initialization, + * specialized for sample types larger than 8-bit. + * This variant initializes the decode operators inside the kernel from level bounds. + * + * @param 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 temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param d_samples + * The pointer to the input sequence of sample items. + * The samples from different channels are assumed to be interleaved + * (e.g., an array of 32-bit pixels where each pixel consists of four RGBA 8-bit samples). + * + * @param d_output_histograms + * The pointers to the histogram counter output arrays, one for each active channel. + * For channeli, the allocation length of `d_histograms[i]` should be + * `num_output_levels[i] - 1`. + * + * @param num_output_levels + * The number of bin level boundaries for delineating histogram samples in each active channel. + * Implies that the number of bins for channeli is + * `num_output_levels[i] - 1`. + * + * @param lower_level + * The lower sample value bound (inclusive) for the lowest histogram bin in each active channel. + * + * @param upper_level + * The upper sample value bound (exclusive) for the highest histogram bin in each active + * channel. + * + * @param num_row_pixels + * The number of multi-channel pixels per row in the region of interest + * + * @param num_rows + * The number of rows in the region of interest + * + * @param row_stride_samples + * The number of samples between starts of consecutive rows in the region of interest + * + * @param stream + * CUDA stream to launch kernels within. Default is stream0. + * + */ +template < + int NUM_CHANNELS, + int NUM_ACTIVE_CHANNELS, + typename SampleIteratorT, + typename CounterT, + typename LevelT, + typename OffsetT, + typename PolicySelector, + typename SampleT = it_value_t, /// The sample value type of the input iterator + typename KernelSource = + DeviceHistogramKernelSource, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY, + typename LowerLevelArrayT = ::cuda::std::array, + typename UpperLevelArrayT = ::cuda::std::array> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t __dispatch_even_device_init( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_output_histograms, + ::cuda::std::array num_output_levels, + LowerLevelArrayT lower_level, + UpperLevelArrayT upper_level, + OffsetT num_row_pixels, + OffsetT num_rows, + OffsetT row_stride_samples, + cudaStream_t stream, + ::cuda::std::false_type /*is_byte_sample*/, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + int max_levels = num_output_levels[0]; + + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + int num_levels = num_output_levels[channel]; + if (kernel_source.MayOverflow(num_levels - 1, upper_level, lower_level, channel)) + { + // Make sure to also return a reasonable value for `temp_storage_bytes` in case of + // an overflow of the bin computation, in which case a subsequent algorithm + // invocation will also fail + if (!d_temp_storage) + { + temp_storage_bytes = 1U; + } + return cudaErrorInvalidValue; + } + + if (num_levels > max_levels) + { + max_levels = num_levels; + } + } + int max_num_output_bins = max_levels - 1; + + if (max_num_output_bins > detail::histogram::max_privatized_smem_bins) + { + // Dispatch shared-privatized approach + constexpr int PRIVATIZED_SMEM_BINS = 0; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_output_levels, + num_output_levels, + upper_level, + lower_level, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + } + else + { + // Dispatch shared-privatized approach + constexpr int PRIVATIZED_SMEM_BINS = detail::histogram::max_privatized_smem_bins; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_output_levels, + num_output_levels, + upper_level, + lower_level, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + } + + return cudaSuccess; +} + +/** + * Dispatch routine for HistogramEven with device-side decode operator initialization, + * specialized for 8-bit sample types + * (computes 256-bin privatized histograms and then reduces to user-specified levels). + * This variant initializes the decode operators inside the kernel from level bounds. + * + * @param 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 temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param d_samples + * The pointer to the input sequence of sample items. The samples from different channels are + * assumed to be interleaved (e.g., an array of 32-bit pixels where each pixel consists of + * four RGBA 8-bit samples). + * + * @param d_output_histograms + * The pointers to the histogram counter output arrays, one for each active channel. + * For channeli, the allocation length of `d_histograms[i]` should be + * `num_output_levels[i] - 1`. + * + * @param num_output_levels + * The number of bin level boundaries for delineating histogram samples in each active channel. + * Implies that the number of bins for channeli is + * `num_output_levels[i] - 1`. + * + * @param lower_level + * The lower sample value bound (inclusive) for the lowest histogram bin in each active channel. + * + * @param upper_level + * The upper sample value bound (exclusive) for the highest histogram bin in each active + * channel. + * + * @param num_row_pixels + * The number of multi-channel pixels per row in the region of interest + * + * @param num_rows + * The number of rows in the region of interest + * + * @param row_stride_samples + * The number of samples between starts of consecutive rows in the region of interest + * + * @param stream + * CUDA stream to launch kernels within. Default is stream0. + * + */ +template < + int NUM_CHANNELS, + int NUM_ACTIVE_CHANNELS, + typename SampleIteratorT, + typename CounterT, + typename LevelT, + typename OffsetT, + typename PolicySelector, + typename SampleT = it_value_t, /// The sample value type of the input iterator + typename KernelSource = + DeviceHistogramKernelSource, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY, + typename LowerLevelArrayT = ::cuda::std::array, + typename UpperLevelArrayT = ::cuda::std::array> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t __dispatch_even_device_init( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_output_histograms, + ::cuda::std::array num_output_levels, + LowerLevelArrayT lower_level, + UpperLevelArrayT upper_level, + OffsetT num_row_pixels, + OffsetT num_rows, + OffsetT row_stride_samples, + cudaStream_t stream, + ::cuda::std::true_type /*is_byte_sample*/, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + ::cuda::std::array num_privatized_levels; + int max_levels = num_output_levels[0]; + + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + num_privatized_levels[channel] = 257; + + int num_levels = num_output_levels[channel]; + if (kernel_source.MayOverflow(num_levels - 1, upper_level, lower_level, channel)) + { + // Make sure to also return a reasonable value for `temp_storage_bytes` in case of + // an overflow of the bin computation, in which case a subsequent algorithm + // invocation will also fail + if (!d_temp_storage) + { + temp_storage_bytes = 1U; + } + return cudaErrorInvalidValue; + } + + if (num_levels > max_levels) + { + max_levels = num_levels; + } + } + int max_num_output_bins = max_levels - 1; + + constexpr int PRIVATIZED_SMEM_BINS = 256; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_privatized_levels, + num_output_levels, + upper_level, + lower_level, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + + return cudaSuccess; +} + +// TODO(bgruber): drop in CCCL 4.0 +template +_CCCL_HOST_DEVICE_API constexpr auto convert_pdl_trigger(int) + -> decltype(ActivePolicy::init_kernel_pdl_trigger_max_bins) +{ + return ActivePolicy::init_kernel_pdl_trigger_max_bins; +} + +// TODO(bgruber): drop in CCCL 4.0 +template +_CCCL_HOST_DEVICE_API constexpr auto convert_pdl_trigger(long) +{ + return 0; +} + +// TODO(bgruber): drop in CCCL 4.0 +template +_CCCL_HOST_DEVICE_API constexpr auto convert_policy() -> HistogramPolicy +{ + using ap = typename ActivePolicy::AgentHistogramPolicyT; + return HistogramPolicy{ + ap::BLOCK_THREADS, + ap::PIXELS_PER_THREAD, + ap::VEC_SIZE, + ap::LOAD_ALGORITHM, + ap::LOAD_MODIFIER, + ap::IS_RLE_COMPRESS, + ap::MEM_PREFERENCE, + ap::IS_WORK_STEALING, + convert_pdl_trigger(0)}; +} + +// TODO(bgruber): drop in CCCL 4.0 +template +struct policy_selector_from_max_policy +{ +private: + struct extract_policy_dispatch_t + { + HistogramPolicy& policy; + + template + _CCCL_HOST_DEVICE_API constexpr cudaError_t Invoke() + { + policy = convert_policy(); + return cudaSuccess; + } + }; + +public: + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const -> HistogramPolicy + { + NV_IF_ELSE_TARGET(NV_IS_HOST, + ({ + HistogramPolicy policy{}; + extract_policy_dispatch_t dispatch{policy}; + _CCCL_VERIFY(MaxPolicy::Invoke(cc.get() * 10, dispatch) == cudaSuccess, ""); + return policy; + }), + ({ return convert_policy(); })); + } +}; + +template < + int NUM_CHANNELS, + int NUM_ACTIVE_CHANNELS, + typename SampleIteratorT, + typename CounterT, + typename LevelT, + typename OffsetT, + bool IsByteSample, + typename PolicySelector, + typename SampleT = it_value_t, /// The sample value type of the input iterator + typename KernelSource = + DeviceHistogramKernelSource, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +CUB_RUNTIME_FUNCTION static cudaError_t dispatch_range( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_output_histograms, + ::cuda::std::array num_output_levels, + ::cuda::std::array d_levels, + OffsetT num_row_pixels, + OffsetT num_rows, + OffsetT row_stride_samples, + cudaStream_t stream, + ::cuda::std::bool_constant, + PolicySelector policy_selector, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + if constexpr (IsByteSample) + { + using TransformsT = Transforms; + + // Use the pass-thru transform op for converting samples to privatized bins + using PrivatizedDecodeOpT = typename TransformsT::PassThruTransform; + + // Use the search transform op for converting privatized bins to output bins + using OutputDecodeOpT = typename TransformsT::template SearchTransform; + + ::cuda::std::array num_privatized_levels; + ::cuda::std::array privatized_decode_op{}; + ::cuda::std::array output_decode_op{}; + int max_levels = num_output_levels[0]; + + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + num_privatized_levels[channel] = 257; + output_decode_op[channel].Init(d_levels[channel], num_output_levels[channel]); + + if (num_output_levels[channel] > max_levels) + { + max_levels = num_output_levels[channel]; + } + } + int max_num_output_bins = max_levels - 1; + + constexpr int PRIVATIZED_SMEM_BINS = 256; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_privatized_levels, + num_output_levels, + output_decode_op, + privatized_decode_op, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + } + else + { + using TransformsT = Transforms; + + // Use the search transform op for converting samples to privatized bins + using PrivatizedDecodeOpT = typename TransformsT::template SearchTransform; + + // Use the pass-thru transform op for converting privatized bins to output bins + using OutputDecodeOpT = typename TransformsT::PassThruTransform; + + ::cuda::std::array privatized_decode_op{}; + ::cuda::std::array output_decode_op{}; + int max_levels = num_output_levels[0]; + + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + privatized_decode_op[channel].Init(d_levels[channel], num_output_levels[channel]); + if (num_output_levels[channel] > max_levels) + { + max_levels = num_output_levels[channel]; + } + } + int max_num_output_bins = max_levels - 1; + + // Dispatch + if (max_num_output_bins > max_privatized_smem_bins) + { + // Too many bins to keep in shared memory. + constexpr int PRIVATIZED_SMEM_BINS = 0; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_output_levels, + num_output_levels, + output_decode_op, + privatized_decode_op, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + } + else + { + // Dispatch shared-privatized approach + constexpr int PRIVATIZED_SMEM_BINS = max_privatized_smem_bins; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_output_levels, + num_output_levels, + output_decode_op, + privatized_decode_op, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + } + } + + return cudaSuccess; +} + +template < + int NUM_CHANNELS, + int NUM_ACTIVE_CHANNELS, + typename SampleIteratorT, + typename CounterT, + typename LevelT, + typename OffsetT, + bool IsByteSample, + typename PolicySelector, + typename SampleT = it_value_t, /// The sample value type of the input iterator + typename KernelSource = + DeviceHistogramKernelSource, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t dispatch_even( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_output_histograms, + ::cuda::std::array num_output_levels, + ::cuda::std::array lower_level, + ::cuda::std::array upper_level, + OffsetT num_row_pixels, + OffsetT num_rows, + OffsetT row_stride_samples, + cudaStream_t stream, + ::cuda::std::bool_constant, + PolicySelector policy_selector, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + if constexpr (IsByteSample) + { + using TransformsT = Transforms; + + // Use the pass-thru transform op for converting samples to privatized bins + using PrivatizedDecodeOpT = typename TransformsT::PassThruTransform; + + // Use the scale transform op for converting privatized bins to output bins + using OutputDecodeOpT = typename TransformsT::ScaleTransform; + + using CommonT = typename TransformsT::ScaleTransform::CommonT; + + ::cuda::std::array num_privatized_levels; + ::cuda::std::array privatized_decode_op{}; + ::cuda::std::array output_decode_op{}; + int max_levels = num_output_levels[0]; + + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + num_privatized_levels[channel] = 257; + + int num_levels = num_output_levels[channel]; + if (kernel_source.MayOverflow(static_cast(num_levels - 1), upper_level, lower_level, channel)) + { + if (!d_temp_storage) + { + temp_storage_bytes = 1U; + } + return cudaErrorInvalidValue; + } + + output_decode_op[channel].Init(num_levels, upper_level[channel], lower_level[channel]); + + if (num_levels > max_levels) + { + max_levels = num_levels; + } + } + int max_num_output_bins = max_levels - 1; + + constexpr int PRIVATIZED_SMEM_BINS = 256; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_privatized_levels, + num_output_levels, + output_decode_op, + privatized_decode_op, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + } + else + { + using TransformsT = Transforms; + + // Use the scale transform op for converting samples to privatized bins + using PrivatizedDecodeOpT = typename TransformsT::ScaleTransform; + + // Use the pass-thru transform op for converting privatized bins to output bins + using OutputDecodeOpT = typename TransformsT::PassThruTransform; + + using CommonT = typename TransformsT::ScaleTransform::CommonT; + + ::cuda::std::array privatized_decode_op{}; + ::cuda::std::array output_decode_op{}; + int max_levels = num_output_levels[0]; + + for (int channel = 0; channel < NUM_ACTIVE_CHANNELS; ++channel) + { + int num_levels = num_output_levels[channel]; + if (kernel_source.MayOverflow(static_cast(num_levels - 1), upper_level, lower_level, channel)) + { + if (!d_temp_storage) + { + temp_storage_bytes = 1U; + } + return cudaErrorInvalidValue; + } + + privatized_decode_op[channel].Init(num_levels, upper_level[channel], lower_level[channel]); + + if (num_levels > max_levels) + { + max_levels = num_levels; + } + } + int max_num_output_bins = max_levels - 1; + + if (max_num_output_bins > max_privatized_smem_bins) + { + constexpr int PRIVATIZED_SMEM_BINS = 0; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_output_levels, + num_output_levels, + output_decode_op, + privatized_decode_op, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + } + else + { + constexpr int PRIVATIZED_SMEM_BINS = max_privatized_smem_bins; + + if (const auto error = CubDebug( + (detail::histogram::dispatch( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_output_levels, + num_output_levels, + output_decode_op, + privatized_decode_op, + max_num_output_bins, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + policy_selector, + kernel_source, + launcher_factory)))) + { + return error; + } + } + } + + return cudaSuccess; +} +} // namespace detail::histogram + +/****************************************************************************** + * Dispatch + ******************************************************************************/ + +// TODO(bgruber): remove in CCCL 4.0 +/** + * Utility class for dispatching the appropriately-tuned kernels for DeviceHistogram + * + * Deprecated [Since 3.5] + * + * @tparam NUM_CHANNELS + * Number of channels interleaved in the input data (may be greater than the number of channels + * being actively histogrammed) + * + * @tparam NUM_ACTIVE_CHANNELS + * Number of channels actively being histogrammed + * + * @tparam SampleIteratorT + * Random-access input iterator type for reading input items @iterator + * + * @tparam CounterT + * Integer type for counting sample occurrences per histogram bin + * + * @tparam LevelT + * Type for specifying bin level boundaries + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam PolicyHub + * Implementation detail, do not specify directly, requirements on the + * content of this type are subject to breaking change. + */ +template < + int NUM_CHANNELS, + int NUM_ACTIVE_CHANNELS, + typename SampleIteratorT, + typename CounterT, + typename LevelT, + typename OffsetT, + typename PolicyHub = void, // if user passes a custom Policy this should not be void + typename SampleT = cub::detail::it_value_t, /// The sample value type of the input iterator + typename KernelSource = detail::histogram:: + DeviceHistogramKernelSource, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceHistogram") DispatchHistogram +{ + static_assert(NUM_CHANNELS <= 4, "Histograms only support up to 4 channels"); + static_assert(NUM_ACTIVE_CHANNELS <= NUM_CHANNELS, + "Active channels must be at most the number of total channels of the input samples"); + + //--------------------------------------------------------------------- + // Dispatch entrypoints + //--------------------------------------------------------------------- + + //--------------------------------------------------------------------- + // Default (host-init) dispatch entrypoints + // These methods initialize decode operators on the host before kernel launch. + //--------------------------------------------------------------------- + + /** + * Dispatch routine for HistogramRange with host-side decode operator initialization. + * This variant initializes the decode operators on the host before kernel launch. + * + * @param 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 temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param d_samples + * The pointer to the multi-channel input sequence of data samples. + * The samples from different channels are assumed to be interleaved + * (e.g., an array of 32-bit pixels where each pixel consists of four RGBA 8-bit samples). + * + * @param d_output_histograms + * The pointers to the histogram counter output arrays, one for each active channel. + * For channeli, the allocation length of `d_histograms[i]` should be + * `num_output_levels[i] - 1`. + * + * @param num_output_levels + * The number of boundaries (levels) for delineating histogram samples in each active channel. + * Implies that the number of bins for channeli is + * `num_output_levels[i] - 1`. + * + * @param d_levels + * The pointers to the arrays of boundaries (levels), one for each active channel. + * Bin ranges are defined by consecutive boundary pairings: lower sample value boundaries are + * inclusive and upper sample value boundaries are exclusive. + * + * @param num_row_pixels + * The number of multi-channel pixels per row in the region of interest + * + * @param num_rows + * The number of rows in the region of interest + * + * @param row_stride_samples + * The number of samples between starts of consecutive rows in the region of interest + * + * @param stream + * CUDA stream to launch kernels within. Default is stream0. + */ + template , + /* fallback_policy_hub */ + detail::histogram::policy_hub, + PolicyHub>::MaxPolicy, + bool IsByteSample> + CUB_RUNTIME_FUNCTION static cudaError_t DispatchRange( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_output_histograms, + ::cuda::std::array num_output_levels, + ::cuda::std::array d_levels, + OffsetT num_row_pixels, + OffsetT num_rows, + OffsetT row_stride_samples, + cudaStream_t stream, + ::cuda::std::bool_constant is_byte_sample, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + [[maybe_unused]] MaxPolicyT max_policy = {}) + { + return detail::histogram::dispatch_range( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_output_levels, + d_levels, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + is_byte_sample, + detail::histogram::policy_selector_from_max_policy{}, + kernel_source, + launcher_factory); + } + + /** + * Dispatch routine for HistogramEven with host-side decode operator initialization. + * This variant initializes the decode operators on the host before kernel launch. + * + * @param 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 temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param d_samples + * The pointer to the input sequence of sample items. + * The samples from different channels are assumed to be interleaved + * (e.g., an array of 32-bit pixels where each pixel consists of four RGBA 8-bit samples). + * + * @param d_output_histograms + * The pointers to the histogram counter output arrays, one for each active channel. + * For channeli, the allocation length of `d_histograms[i]` should be + * `num_output_levels[i] - 1`. + * + * @param num_output_levels + * The number of bin level boundaries for delineating histogram samples in each active channel. + * Implies that the number of bins for channeli is + * `num_output_levels[i] - 1`. + * + * @param lower_level + * The lower sample value bound (inclusive) for the lowest histogram bin in each active channel. + * + * @param upper_level + * The upper sample value bound (exclusive) for the highest histogram bin in each active + * channel. + * + * @param num_row_pixels + * The number of multi-channel pixels per row in the region of interest + * + * @param num_rows + * The number of rows in the region of interest + * + * @param row_stride_samples + * The number of samples between starts of consecutive rows in the region of interest + * + * @param stream + * CUDA stream to launch kernels within. Default is stream0. + * + */ + template , + /* fallback_policy_hub */ + detail::histogram::policy_hub, + PolicyHub>::MaxPolicy, + bool IsByteSample> + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t DispatchEven( + void* d_temp_storage, + size_t& temp_storage_bytes, + SampleIteratorT d_samples, + ::cuda::std::array d_output_histograms, + ::cuda::std::array num_output_levels, + ::cuda::std::array lower_level, + ::cuda::std::array upper_level, + OffsetT num_row_pixels, + OffsetT num_rows, + OffsetT row_stride_samples, + cudaStream_t stream, + ::cuda::std::bool_constant is_byte_sample, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + [[maybe_unused]] MaxPolicyT max_policy = {}) + { + return detail::histogram::dispatch_even( + d_temp_storage, + temp_storage_bytes, + d_samples, + d_output_histograms, + num_output_levels, + lower_level, + upper_level, + num_row_pixels, + num_rows, + row_stride_samples, + stream, + is_byte_sample, + detail::histogram::policy_selector_from_max_policy{}, + kernel_source, + launcher_factory); + } +}; + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_merge.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_merge.cuh new file mode 100644 index 00000000..d4054f8a --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_merge.cuh @@ -0,0 +1,336 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +CUB_NAMESPACE_BEGIN +namespace detail::merge +{ +inline constexpr int fallback_BLOCK_THREADS = 64; +inline constexpr int fallback_ITEMS_PER_THREAD = 1; + +// TODO(bgruber): we should choose the MergePolicy rather than the agent, but before C++20 this is more verbose +template +class choose_merge_agent +{ + static constexpr MergePolicy active_policy = PolicyGetter{}(); + + using default_load2sh_agent_t = + agent_t; + using default_noload2sh_agent_t = + agent_t; + + using fallback_agent_t = + agent_t; + + static constexpr bool use_default_load2sh = + sizeof(typename default_load2sh_agent_t::TempStorage) <= max_smem_per_block; + // Use fallback if merge agent exceeds maximum shared memory, but the fallback agent still fits, else use + // vsmem-compatible version, so noload2sh + static constexpr bool use_fallback = sizeof(typename fallback_agent_t::TempStorage) <= max_smem_per_block; + +public: + using type = + ::cuda::std::conditional_t>; +}; + +// Computes the merge path intersections at equally wide intervals. The approach is outlined in the paper: +// Odeh et al, "Merge Path - Parallel Merging Made Simple" * doi : 10.1109 / IPDPSW .2012.202 +// The algorithm is the same as AgentPartition for merge sort, but that agent handles a lot more. +template +_CCCL_KERNEL_ATTRIBUTES void device_partition_merge_path_kernel( + const KeyIt1 keys1, + const Offset keys1_count, + const KeyIt2 keys2, + const Offset keys2_count, + const Offset num_diagonals, + Offset* key1_beg_offsets, + CompareOp compare_op) +{ + // items_per_tile must be the same of the merge kernel later, so we have to consider whether a fallback agent will be + // selected for the merge agent that changes the tile size + constexpr int items_per_tile = + choose_merge_agent, + KeyIt1, + ValueIt1, + KeyIt2, + ValueIt2, + KeyIt3, + ValueIt3, + Offset, + CompareOp>::type::items_per_tile; + const Offset diagonal_idx = + static_cast(blockDim.x * blockIdx.x + threadIdx.x); // NOLINT(bugprone-misplaced-widening-cast) + if (diagonal_idx < num_diagonals) + { + const Offset diagonal_num = (::cuda::std::min) (diagonal_idx * items_per_tile, keys1_count + keys2_count); + key1_beg_offsets[diagonal_idx] = cub::MergePath(keys1, keys2, keys1_count, keys2_count, diagonal_num, compare_op); + } +} + +template +__launch_bounds__( + choose_merge_agent, + KeyIt1, + ValueIt1, + KeyIt2, + ValueIt2, + KeyIt3, + ValueIt3, + Offset, + CompareOp>::type::threads_per_block) + _CCCL_KERNEL_ATTRIBUTES void device_merge_kernel( + const KeyIt1 keys1, + const ValueIt1 items1, + const Offset num_keys1, + const KeyIt2 keys2, + const ValueIt2 items2, + const Offset num_keys2, + const KeyIt3 keys_result, + const ValueIt3 items_result, + CompareOp compare_op, + Offset* key1_beg_offsets, + vsmem_t global_temp_storage) +{ + using key_t = it_value_t; + static_assert(::cuda::std::is_invocable_v, "Comparison operator cannot compare two keys"); + static_assert(::cuda::std::is_convertible_v<::cuda::std::invoke_result_t, bool>, + "Comparison operator must be convertible to bool"); + + using MergeAgent = typename choose_merge_agent< + device_policy_getter, + KeyIt1, + ValueIt1, + KeyIt2, + ValueIt2, + KeyIt3, + ValueIt3, + Offset, + CompareOp>::type; + + using vsmem_helper_t = vsmem_helper_impl; + __shared__ typename vsmem_helper_t::static_temp_storage_t shared_temp_storage; + auto& temp_storage = vsmem_helper_t::get_temp_storage(shared_temp_storage, global_temp_storage); + MergeAgent{ + temp_storage.Alias(), + keys1, + items1, + num_keys1, + keys2, + items2, + num_keys2, + keys_result, + items_result, + compare_op, + key1_beg_offsets}(); + vsmem_helper_t::discard_temp_storage(temp_storage); +} + +template , + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires merge_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyIt1 d_keys1, + ValueIt1 d_values1, + Offset num_items1, + KeyIt2 d_keys2, + ValueIt2 d_values2, + Offset num_items2, + KeyIt3 d_keys_out, + ValueIt3 d_values_out, + CompareOp compare_op, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelLauncherFactory launcher_factory = {}) +{ + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + return dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) { +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << policy_getter(); + _CubLog("Dispatching DeviceMerge to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + static_assert(::cuda::std::is_empty_v); + using AgentT = typename choose_merge_agent< + decltype(policy_getter), + KeyIt1, + ValueIt1, + KeyIt2, + ValueIt2, + KeyIt3, + ValueIt3, + Offset, + CompareOp>::type; + + const auto num_tiles = ::cuda::ceil_div(num_items1 + num_items2, AgentT::items_per_tile); + void* allocations[2] = {nullptr, nullptr}; + { + const size_t key1_beg_offsets_size = (1 + num_tiles) * sizeof(Offset); + const size_t virtual_shared_memory_size = num_tiles * vsmem_helper_impl::vsmem_per_block; + const size_t allocation_sizes[2] = {key1_beg_offsets_size, virtual_shared_memory_size}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + } + + if (d_temp_storage == nullptr || num_tiles == 0) + { + return cudaSuccess; + } + + auto key1_beg_offsets = static_cast(allocations[0]); + + // merge path kernel + { + const Offset num_diagonals = num_tiles + 1; + constexpr int threads_per_partition_block = 256; + const int partition_grid_size = static_cast(::cuda::ceil_div(num_diagonals, threads_per_partition_block)); + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + partition_grid_size, threads_per_partition_block, 0, stream) + .doit(device_partition_merge_path_kernel< + PolicySelector, + KeyIt1, + ValueIt1, + KeyIt2, + ValueIt2, + KeyIt3, + ValueIt3, + Offset, + CompareOp>, + d_keys1, + num_items1, + d_keys2, + num_items2, + num_diagonals, + key1_beg_offsets, + compare_op))) + { + return error; + } + if (const auto error = CubDebug(DebugSyncStream(stream))) + { + return error; + } + } + + // merge kernel + { + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron( + static_cast(num_tiles), static_cast(AgentT::threads_per_block), 0, stream) + .doit( + device_merge_kernel, + d_keys1, + d_values1, + num_items1, + d_keys2, + d_values2, + num_items2, + d_keys_out, + d_values_out, + compare_op, + key1_beg_offsets, + vsmem_t{allocations[1]}))) + { + return error; + } + if (const auto error = CubDebug(DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; + }); +} +} // namespace detail::merge +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_merge_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_merge_sort.cuh new file mode 100644 index 00000000..b8ac3fa6 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_merge_sort.cuh @@ -0,0 +1,668 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::merge_sort +{ +template +struct DeviceMergeSortKernelSource +{ +#if _CCCL_HAS_CONCEPTS() + static_assert(detail::merge_sort::merge_sort_policy_selector); +#endif // _CCCL_HAS_CONCEPTS() + + using KeyT = cub::detail::it_value_t; + using ValueT = cub::detail::it_value_t; + + CUB_DEFINE_KERNEL_GETTER( + MergeSortBlockSortKernel, + DeviceMergeSortBlockSortKernel< + PolicySelectorT, + KeyInputIteratorT, + ValueInputIteratorT, + KeyIteratorT, + ValueIteratorT, + OffsetT, + CompareOpT, + KeyT, + ValueT>); + + CUB_DEFINE_KERNEL_GETTER(MergeSortPartitionKernel, + DeviceMergeSortPartitionKernel); + + CUB_DEFINE_KERNEL_GETTER( + MergeSortMergeKernel, + DeviceMergeSortMergeKernel); + + CUB_RUNTIME_FUNCTION static constexpr size_t KeySize() + { + return sizeof(KeyT); + } + + CUB_RUNTIME_FUNCTION static constexpr size_t ValueSize() + { + return sizeof(ValueT); + } +}; +} // namespace detail::merge_sort + +/******************************************************************************* + * Policy + ******************************************************************************/ + +// TODO(bgruber): remove in CCCL 4.0 +//! Deprecated [Since 3.5] +template , + typename KernelSource = detail::merge_sort::DeviceMergeSortKernelSource< + detail::merge_sort::policy_selector_from_hub, + KeyInputIteratorT, + ValueInputIteratorT, + KeyIteratorT, + ValueIteratorT, + OffsetT, + CompareOpT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY, + typename KeyT = cub::detail::it_value_t, + typename ValueT = cub::detail::it_value_t> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceMergeSort") DispatchMergeSort +{ + /// Whether or not there are values to be trucked along with keys + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + + // Problem state + + /// Device-accessible allocation of temporary storage. When nullptr, the required + /// allocation size is written to \p temp_storage_bytes and no work is done. + void* d_temp_storage; + + /// Reference to size in bytes of \p d_temp_storage allocation + size_t& temp_storage_bytes; + + /// Pointer to the input sequence of unsorted input keys + KeyInputIteratorT d_input_keys; + + /// Pointer to the input sequence of unsorted input values + ValueInputIteratorT d_input_items; + + /// Pointer to the output sequence of sorted input keys + KeyIteratorT d_output_keys; + + /// Pointer to the output sequence of sorted input values + ValueIteratorT d_output_items; + + /// Number of items to sort + OffsetT num_items; + + /// Comparison function object which returns true if the first argument is + /// ordered before the second + CompareOpT compare_op; + + /// CUDA stream to launch kernels within. Default is stream0. + cudaStream_t stream; + + int ptx_version; + + KernelSource kernel_source; + + KernelLauncherFactory launcher_factory; + + // Constructor + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchMergeSort( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_input_keys, + ValueInputIteratorT d_input_items, + KeyIteratorT d_output_keys, + ValueIteratorT d_output_items, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream, + int ptx_version, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_input_keys(d_input_keys) + , d_input_items(d_input_items) + , d_output_keys(d_output_keys) + , d_output_items(d_output_items) + , num_items(num_items) + , compare_op(compare_op) + , stream(stream) + , ptx_version(ptx_version) + , kernel_source(kernel_source) + , launcher_factory(launcher_factory) + {} + +private: + template + struct policy_getter + { + _CCCL_HOST_DEVICE_API constexpr auto operator()() -> MergeSortPolicy + { + using mp = typename ActivePolicyT::MergeSortPolicy; + return {mp::BLOCK_THREADS, mp::ITEMS_PER_THREAD, mp::LOAD_ALGORITHM, mp::LOAD_MODIFIER, mp::STORE_ALGORITHM}; + } + }; + +public: + // Invocation + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke([[maybe_unused]] ActivePolicyT = {}) + { + if (num_items == 0) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + static constexpr auto policy = detail::merge_sort::merge_sort_vsmem_helper_t< + policy_getter, + KeyInputIteratorT, + ValueInputIteratorT, + KeyIteratorT, + ValueIteratorT, + OffsetT, + CompareOpT, + KeyT, + ValueT>::policy; + static_assert(1 <= policy.threads_per_block && policy.threads_per_block <= 1024, + "Number of threads per block need to be inside [1;1024]"); + static_assert(1 <= policy.items_per_thread, "Number of items per thread needs to be at least 1"); + constexpr auto tile_size = policy.threads_per_block * policy.items_per_thread; + const auto num_tiles = ::cuda::ceil_div(num_items, tile_size); + + const auto merge_partitions_size = static_cast(1 + num_tiles) * sizeof(OffsetT); + const auto temporary_keys_storage_size = static_cast(num_items * kernel_source.KeySize()); + const auto temporary_values_storage_size = static_cast(num_items * kernel_source.ValueSize()) * !KEYS_ONLY; + + /** + * Merge sort supports large types, which can lead to excessive shared memory size requirements. In these cases, + * merge sort allocates virtual shared memory that resides in global memory. + */ + + using merge_sort_vsmem_t = detail::merge_sort::merge_sort_vsmem_helper_t< + policy_getter, + KeyInputIteratorT, + ValueInputIteratorT, + KeyIteratorT, + ValueIteratorT, + OffsetT, + CompareOpT, + KeyT, + ValueT>; + const ::cuda::std::size_t block_sort_smem_size = + num_tiles * detail::vsmem_helper_impl::vsmem_per_block; + const ::cuda::std::size_t merge_smem_size = + num_tiles * detail::vsmem_helper_impl::vsmem_per_block; + const ::cuda::std::size_t virtual_shared_memory_size = (::cuda::std::max) (block_sort_smem_size, merge_smem_size); + + void* allocations[4] = {nullptr, nullptr, nullptr, nullptr}; + size_t allocation_sizes[4] = { + merge_partitions_size, temporary_keys_storage_size, temporary_values_storage_size, virtual_shared_memory_size}; + + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage allocation + return cudaSuccess; + } + + const int num_passes = ::cuda::ceil_ilog2(num_tiles); + + /* + * The algorithm consists of stages. At each stage, there are input and output arrays. There are two pairs of + * arrays allocated (keys and items). One pair is from function arguments and another from temporary storage. Ping + * is a helper variable that controls which of these two pairs of arrays is an input and which is an output for a + * current stage. If the ping is true - the current stage stores its result in the temporary storage. The + * temporary storage acts as input data otherwise. + * + * Block sort is executed before the main loop. It stores its result in the pair of arrays that will be an input + * of the next stage. The initial value of the ping variable is selected so that the result of the final stage is + * stored in the input arrays. + */ + bool ping = num_passes % 2 == 0; + + auto merge_partitions = static_cast(allocations[0]); + auto keys_buffer = static_cast(allocations[1]); + auto items_buffer = static_cast(allocations[2]); + + const int threads_per_block = policy.threads_per_block; + + // Invoke DeviceMergeSortBlockSortKernel + launcher_factory( + static_cast(num_tiles), threads_per_block, 0, stream, /* dependent launch */ ptx_version >= 900) + .doit(kernel_source.MergeSortBlockSortKernel(), + ping, + d_input_keys, + d_input_items, + d_output_keys, + d_output_items, + num_items, + keys_buffer, + items_buffer, + compare_op, + cub::detail::vsmem_t{allocations[3]}); + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + const OffsetT num_partitions = num_tiles + 1; + constexpr int threads_per_partition_block = 256; + const int partition_grid_size = static_cast(::cuda::ceil_div(num_partitions, threads_per_partition_block)); + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + for (int pass = 0; pass < num_passes; ++pass, ping = !ping) + { + const OffsetT target_merged_tiles_number = OffsetT(2) << pass; + + // Partition + launcher_factory( + partition_grid_size, threads_per_partition_block, 0, stream, /* dependent launch */ ptx_version >= 900) + .doit(kernel_source.MergeSortPartitionKernel(), + ping, + d_output_keys, + keys_buffer, + num_items, + num_partitions, + merge_partitions, + compare_op, + target_merged_tiles_number, + tile_size); + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Merge + launcher_factory( + static_cast(num_tiles), threads_per_block, 0, stream, /* dependent launch */ ptx_version >= 900) + .doit(kernel_source.MergeSortMergeKernel(), + ping, + d_output_keys, + d_output_items, + num_items, + keys_buffer, + items_buffer, + compare_op, + merge_partitions, + target_merged_tiles_number, + cub::detail::vsmem_t{allocations[3]}); + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + } + + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_input_keys, + ValueInputIteratorT d_input_items, + KeyIteratorT d_output_keys, + ValueIteratorT d_output_items, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) + { + // Get PTX version + int ptx_version = 0; + if (const auto error = CubDebug(launcher_factory.PtxVersion(ptx_version))) + { + return error; + } + + // Create dispatch functor + DispatchMergeSort dispatch( + d_temp_storage, + temp_storage_bytes, + d_input_keys, + d_input_items, + d_output_keys, + d_output_items, + num_items, + compare_op, + stream, + ptx_version, + kernel_source, + launcher_factory); + + // Dispatch to chained policy + if (const auto error = CubDebug(max_policy.Invoke(ptx_version, dispatch))) + { + return error; + } + + return cudaSuccess; + } +}; + +namespace detail::merge_sort +{ +template , + typename KernelSource = DeviceMergeSortKernelSource, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY, + typename KeyT = it_value_t, + typename ValueT = it_value_t> +#if _CCCL_HAS_CONCEPTS() + requires merge_sort_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeyInputIteratorT d_input_keys, + ValueInputIteratorT d_input_items, + KeyIteratorT d_output_keys, + ValueIteratorT d_output_items, + OffsetT num_items, + CompareOpT compare_op, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + KeyT* = nullptr /* for CCCL.C */, + ValueT* = nullptr /* for CCCL.C */) -> cudaError_t +{ + [[maybe_unused]] constexpr bool keys_only = ::cuda::std::is_same_v; + + if (num_items == 0) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + return detail::dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) -> cudaError_t { +#ifdef CUB_DEFINE_RUNTIME_POLICIES + const MergeSortPolicy active_policy = policy_getter(); +#else // CUB_DEFINE_RUNTIME_POLICIES + using vsmem_adapted_agents = merge_sort_vsmem_helper_t< + decltype(policy_getter), + KeyInputIteratorT, + ValueInputIteratorT, + KeyIteratorT, + ValueIteratorT, + OffsetT, + CompareOpT, + KeyT, + ValueT>; + constexpr MergeSortPolicy active_policy = vsmem_adapted_agents::policy; +#endif // CUB_DEFINE_RUNTIME_POLICIES + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceMergeSort to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + _CCCL_ASSERT(1 <= active_policy.threads_per_block && active_policy.threads_per_block <= 1024, + "Number of threads per block need to be inside [1;1024]"); + _CCCL_ASSERT(1 <= active_policy.items_per_thread, "Number of items per thread needs to be at least 1"); + const auto tile_size = active_policy.threads_per_block * active_policy.items_per_thread; + const auto num_tiles = ::cuda::ceil_div(num_items, tile_size); + + const auto merge_partitions_size = static_cast(1 + num_tiles) * sizeof(OffsetT); + const auto temporary_keys_storage_size = static_cast(num_items * kernel_source.KeySize()); + const auto temporary_values_storage_size = static_cast(num_items * kernel_source.ValueSize()) * !keys_only; + +#ifdef CUB_DEFINE_RUNTIME_POLICIES + const ::cuda::std::size_t block_sort_smem_size = 0; + const ::cuda::std::size_t merge_smem_size = 0; +#else // CUB_DEFINE_RUNTIME_POLICIES + const ::cuda::std::size_t block_sort_smem_size = + num_tiles * vsmem_helper_impl::vsmem_per_block; + const ::cuda::std::size_t merge_smem_size = num_tiles * vsmem_helper_impl::vsmem_per_block; +#endif // CUB_DEFINE_RUNTIME_POLICIES + const ::cuda::std::size_t virtual_shared_memory_size = (::cuda::std::max) (block_sort_smem_size, merge_smem_size); + + void* allocations[4] = {nullptr, nullptr, nullptr, nullptr}; + size_t allocation_sizes[4] = { + merge_partitions_size, temporary_keys_storage_size, temporary_values_storage_size, virtual_shared_memory_size}; + + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + const int num_passes = ::cuda::ceil_ilog2(num_tiles); + bool ping = num_passes % 2 == 0; + + auto merge_partitions = static_cast(allocations[0]); + auto keys_buffer = static_cast(allocations[1]); + auto items_buffer = static_cast(allocations[2]); + + if (const auto error = CubDebug( + launcher_factory(static_cast(num_tiles), + active_policy.threads_per_block, + 0, + stream, + /* dependent launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(kernel_source.MergeSortBlockSortKernel(), + ping, + d_input_keys, + d_input_items, + d_output_keys, + d_output_items, + num_items, + keys_buffer, + items_buffer, + compare_op, + cub::detail::vsmem_t{allocations[3]}))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + const OffsetT num_partitions = num_tiles + 1; + constexpr int threads_per_partition_block = 256; + const int partition_grid_size = static_cast(::cuda::ceil_div(num_partitions, threads_per_partition_block)); + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + for (int pass = 0; pass < num_passes; ++pass, ping = !ping) + { + const OffsetT target_merged_tiles_number = OffsetT(2) << pass; + + if (const auto error = CubDebug( + launcher_factory(partition_grid_size, + threads_per_partition_block, + 0, + stream, + /* dependent launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(kernel_source.MergeSortPartitionKernel(), + ping, + d_output_keys, + keys_buffer, + num_items, + num_partitions, + merge_partitions, + compare_op, + target_merged_tiles_number, + tile_size))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + if (const auto error = CubDebug( + launcher_factory(static_cast(num_tiles), + active_policy.threads_per_block, + 0, + stream, + /* dependent launch */ cc >= ::cuda::compute_capability{9, 0}) + .doit(kernel_source.MergeSortMergeKernel(), + ping, + d_output_keys, + d_output_items, + num_items, + keys_buffer, + items_buffer, + compare_op, + merge_partitions, + target_merged_tiles_number, + cub::detail::vsmem_t{allocations[3]}))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + } + + return cudaSuccess; + }); +} +} // namespace detail::merge_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_radix_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_radix_sort.cuh new file mode 100644 index 00000000..06cfc861 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_radix_sort.cuh @@ -0,0 +1,2070 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::DeviceRadixSort provides device-wide, parallel operations for computing a radix sort across + * a sequence of data items residing within device-accessible memory. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +// TODO(bgruber): included for backward compatibility, remove in CCCL 4.0 +#include + +#include + +// suppress warnings triggered by #pragma unroll: +// "warning: loop not unrolled: the optimizer was unable to perform the requested transformation; the transformation +// might be disabled or specified as part of an unsupported transformation ordering [-Wpass-failed=transform-warning]" +_CCCL_DIAG_PUSH +_CCCL_DIAG_SUPPRESS_CLANG("-Wpass-failed") + +CUB_NAMESPACE_BEGIN + +namespace detail::radix_sort +{ +template +struct DeviceRadixSortKernelSource +{ + // PolicySelector must be stateless, so we can pass the type to the kernel + static_assert(::cuda::std::is_empty_v); + + CUB_DEFINE_KERNEL_GETTER(RadixSortSingleTileKernel, + DeviceRadixSortSingleTileKernel); + + CUB_DEFINE_KERNEL_GETTER(RadixSortUpsweepKernel, + DeviceRadixSortUpsweepKernel); + + CUB_DEFINE_KERNEL_GETTER(RadixSortAltUpsweepKernel, + DeviceRadixSortUpsweepKernel); + + CUB_DEFINE_KERNEL_GETTER(DeviceRadixSortScanBinsKernel, RadixSortScanBinsKernel); + + CUB_DEFINE_KERNEL_GETTER( + RadixSortDownsweepKernel, + DeviceRadixSortDownsweepKernel); + + CUB_DEFINE_KERNEL_GETTER( + RadixSortAltDownsweepKernel, + DeviceRadixSortDownsweepKernel); + + CUB_DEFINE_KERNEL_GETTER(RadixSortHistogramKernel, + DeviceRadixSortHistogramKernel); + + CUB_DEFINE_KERNEL_GETTER(RadixSortExclusiveSumKernel, DeviceRadixSortExclusiveSumKernel); + + CUB_DEFINE_KERNEL_GETTER(RadixSortInitBinsAndCountersKernel, DeviceRadixSortInitKernel); + + CUB_DEFINE_KERNEL_GETTER(RadixSortInitLookbackKernel, DeviceRadixSortInitKernel); + + CUB_DEFINE_KERNEL_GETTER( + RadixSortOnesweepKernel, + DeviceRadixSortOnesweepKernel); + + CUB_RUNTIME_FUNCTION static constexpr size_t KeySize() + { + return sizeof(KeyT); + } + + CUB_RUNTIME_FUNCTION static constexpr size_t ValueSize() + { + return sizeof(ValueT); + } + + CUB_RUNTIME_FUNCTION static constexpr KeyT* AdvanceKeys(KeyT* ptr, OffsetT offset) + { + return ptr + offset; + } + + CUB_RUNTIME_FUNCTION static constexpr ValueT* AdvanceValues(ValueT* ptr, OffsetT offset) + { + return ptr + offset; + } +}; +} // namespace detail::radix_sort + +/****************************************************************************** + * Single-problem dispatch + ******************************************************************************/ + +/** + * Utility class for dispatching the appropriately-tuned kernels for device-wide radix sort + * + * Deprecated [Since 3.5] + * + * @tparam SortOrder + * Whether to sort in ascending or descending order + * + * @tparam KeyT + * Key type + * + * @tparam ValueT + * Value type + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam DecomposerT + * Implementation detail, do not specify directly, requirements on the + * content of this type are subject to breaking change. + */ +template , + typename KernelSource = detail::radix_sort::DeviceRadixSortKernelSource< + detail::radix_sort::policy_selector_from_hub, + Order, + KeyT, + ValueT, + OffsetT, + DecomposerT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceRadixSort") DispatchRadixSort +{ + //------------------------------------------------------------------------------ + // Constants + //------------------------------------------------------------------------------ + + // Whether this is a keys-only (or key-value) sort + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + + //------------------------------------------------------------------------------ + // Problem state + //------------------------------------------------------------------------------ + + /// Device-accessible allocation of temporary storage. + // When nullptr, the required allocation size is written to `temp_storage_bytes` and no work is + // done. + void* d_temp_storage; + + /// Reference to size in bytes of `d_temp_storage` allocation + size_t& temp_storage_bytes; + + /// Double-buffer whose current buffer contains the unsorted input keys and, upon return, is + /// updated to point to the sorted output keys + DoubleBuffer& d_keys; + + /// Double-buffer whose current buffer contains the unsorted input values and, upon return, is + /// updated to point to the sorted output values + DoubleBuffer& d_values; + + /// Number of items to sort + OffsetT num_items; + + /// The beginning (least-significant) bit index needed for key comparison + int begin_bit; + + /// The past-the-end (most-significant) bit index needed for key comparison + int end_bit; + + /// CUDA stream to launch kernels within. Default is stream0. + cudaStream_t stream; + + /// PTX version + int ptx_version; + + /// Whether is okay to overwrite source buffers + bool is_overwrite_okay; + + DecomposerT decomposer; + + KernelSource kernel_source; + + KernelLauncherFactory launcher_factory; + + //------------------------------------------------------------------------------ + // Constructor + //------------------------------------------------------------------------------ + + // TODO(bgruber): Remove in CCCL 4.0 + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchRadixSort( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + OffsetT num_items, + int begin_bit, + int end_bit, + bool is_overwrite_okay, + cudaStream_t stream, + int ptx_version, + DecomposerT decomposer = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_keys(d_keys) + , d_values(d_values) + , num_items(num_items) + , begin_bit(begin_bit) + , end_bit(end_bit) + , stream(stream) + , ptx_version(ptx_version) + , is_overwrite_okay(is_overwrite_okay) + , decomposer(decomposer) + , kernel_source(kernel_source) + , launcher_factory(launcher_factory) + {} + + //------------------------------------------------------------------------------ + // Small-problem (single tile) invocation + //------------------------------------------------------------------------------ + + /** + * @brief Invoke a single block to sort in-core + * + * @tparam ActivePolicyT + * Umbrella policy active for the target device + * + * @tparam SingleTileKernelT + * Function type of cub::DeviceRadixSortSingleTileKernel + * + * @param[in] single_tile_kernel + * Kernel function pointer to parameterization of cub::DeviceRadixSortSingleTileKernel + */ + // TODO(bgruber): Remove in CCCL 4.0 + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t + InvokeSingleTile(SingleTileKernelT single_tile_kernel, ActivePolicyT policy = {}) + { + return __invoke_single_tile(single_tile_kernel, detail::radix_sort::convert_policy(policy).single_tile); + } + +private: + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t + __invoke_single_tile(SingleTileKernelT single_tile_kernel, RadixSortDownsweepPolicy policy) + { + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + // Log single_tile_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking single_tile_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, %d SM occupancy, current bit " + "%d, bit_grain %d\n", + 1, + policy.threads_per_block, + (long long) stream, + policy.items_per_thread, + 1, + begin_bit, + policy.radix_bits); +#endif + + // Invoke upsweep_kernel with same grid size as downsweep_kernel + launcher_factory(1, policy.threads_per_block, 0, stream) + .doit(single_tile_kernel, + d_keys.Current(), + d_keys.Alternate(), + d_values.Current(), + d_values.Alternate(), + num_items, + begin_bit, + end_bit, + decomposer); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Update selector + d_keys.selector ^= 1; + d_values.selector ^= 1; + + return cudaSuccess; + } + +public: + //------------------------------------------------------------------------------ + // Normal problem size invocation + //------------------------------------------------------------------------------ + + /** + * Invoke a three-kernel sorting pass at the current bit. + */ + // TODO(bgruber): Remove in CCCL 4.0 + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t InvokePass( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + OffsetT* d_spine, + int /*spine_length*/, + int& current_bit, + PassConfigT& pass_config) + { + int pass_bits = ::cuda::std::min(pass_config.radix_bits, end_bit - current_bit); + +// Log upsweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking upsweep_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, %d SM occupancy, current bit %d, " + "bit_grain %d\n", + pass_config.even_share.grid_size, + pass_config.upsweep_config.threads_per_block, + (long long) stream, + pass_config.upsweep_config.items_per_thread, + pass_config.upsweep_config.sm_occupancy, + current_bit, + pass_bits); +#endif + + // Spine length written by the upsweep kernel in the current pass. + int pass_spine_length = pass_config.even_share.grid_size * pass_config.radix_digits; + + // Invoke upsweep_kernel with same grid size as downsweep_kernel + launcher_factory(pass_config.even_share.grid_size, pass_config.upsweep_config.threads_per_block, 0, stream) + .doit(pass_config.upsweep_kernel, + d_keys_in, + d_spine, + num_items, + current_bit, + pass_bits, + pass_config.even_share, + decomposer); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + +// Log scan_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking scan_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread\n", + 1, + pass_config.scan_config.threads_per_block, + (long long) stream, + pass_config.scan_config.items_per_thread); +#endif + + // Invoke scan_kernel + launcher_factory(1, pass_config.scan_config.threads_per_block, 0, stream) + .doit(pass_config.scan_kernel, d_spine, pass_spine_length); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + +// Log downsweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking downsweep_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, %d SM occupancy\n", + pass_config.even_share.grid_size, + pass_config.downsweep_config.threads_per_block, + (long long) stream, + pass_config.downsweep_config.items_per_thread, + pass_config.downsweep_config.sm_occupancy); +#endif + + // Invoke downsweep_kernel + launcher_factory(pass_config.even_share.grid_size, pass_config.downsweep_config.threads_per_block, 0, stream) + .doit(pass_config.downsweep_kernel, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + d_spine, + num_items, + current_bit, + pass_bits, + pass_config.even_share, + decomposer); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Update current bit + current_bit += pass_bits; + + return cudaSuccess; + } + + // TODO(bgruber): Remove in CCCL 4.0 + /// Pass configuration structure + template + struct PassConfig + { + UpsweepKernelT upsweep_kernel; + detail::KernelConfig upsweep_config; + ScanKernelT scan_kernel; + detail::KernelConfig scan_config; + DownsweepKernelT downsweep_kernel; + detail::KernelConfig downsweep_config; + int radix_bits; + int radix_digits; + int max_downsweep_grid_size; + GridEvenShare even_share; + + // TODO(bgruber): Remove in CCCL 4.0 + /// Initialize pass configuration + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t InitPassConfig( + UpsweepKernelT upsweep_kernel, + ScanKernelT scan_kernel, + DownsweepKernelT downsweep_kernel, + int /*ptx_version*/, + int sm_count, + OffsetT num_items, + ActivePolicyT policy = {}, + UpsweepPolicyT upsweep_policy = {}, + ScanPolicyT scan_policy = {}, + DownsweepPolicyT downsweep_policy = {}, + KernelLauncherFactory launcher_factory = {}) + { + // FIXME(bgruber): we should actually convert upsweep_policy, scan_policy, and downsweep_policy, since they could + // be different from those inside policy. But this is already so far out of any supported scenario that I am + // willing to cut this corner. + const auto p = detail::radix_sort::convert_policy(policy); + __init_pass_config( + upsweep_kernel, + scan_kernel, + downsweep_kernel, + sm_count, + num_items, + p.downsweep.radix_bits, + p.upsweep, + p.scan, + p.downsweep, + launcher_factory); + return __init_pass_config(p); + } + + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t __init_pass_config( + UpsweepKernelT upsweep_kern, + ScanKernelT scan_kern, + DownsweepKernelT downsweep_kern, + int sm_count, + OffsetT num_items, + int pass_radix_bits, + RadixSortUpsweepPolicy upsweep_policy, + ScanPolicy scan_policy, + RadixSortDownsweepPolicy downsweep_policy, + KernelLauncherFactory launcher_factory) + { + this->upsweep_kernel = upsweep_kern; + this->scan_kernel = scan_kern; + this->downsweep_kernel = downsweep_kern; + this->radix_bits = pass_radix_bits; + radix_digits = 1 << radix_bits; + + if (const auto error = CubDebug(upsweep_config.__init(upsweep_kernel, upsweep_policy, launcher_factory))) + { + return error; + } + + if (const auto error = CubDebug(scan_config.__init(scan_kernel, scan_policy.lookback, launcher_factory))) + { + return error; + } + + if (const auto error = CubDebug(downsweep_config.__init(downsweep_kernel, downsweep_policy, launcher_factory))) + { + return error; + } + + max_downsweep_grid_size = (downsweep_config.sm_occupancy * sm_count) * detail::subscription_factor; + + even_share.DispatchInit( + num_items, max_downsweep_grid_size, ::cuda::std::max(downsweep_config.tile_size, upsweep_config.tile_size)); + + return cudaSuccess; + } + }; + + // TODO(bgruber): Remove in CCCL 4.0 + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t InvokeOnesweep(ActivePolicyT policy = {}) + { + return __invoke_onesweep(detail::radix_sort::convert_policy(policy)); + } + +private: + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t __invoke_onesweep(RadixSortPolicy policy) + { + // PortionOffsetT is used for offsets within a portion, and must be signed. + using PortionOffsetT = int; + using AtomicOffsetT = PortionOffsetT; + + // compute temporary storage size + const int RADIX_BITS = policy.onesweep.radix_bits; + const int RADIX_DIGITS = 1 << RADIX_BITS; + const int ONESWEEP_ITEMS_PER_THREAD = policy.onesweep.items_per_thread; + const int ONESWEEP_BLOCK_THREADS = policy.onesweep.threads_per_block; + const int ONESWEEP_TILE_ITEMS = ONESWEEP_ITEMS_PER_THREAD * ONESWEEP_BLOCK_THREADS; + // portions handle inputs with >=2**30 elements, due to the way lookback works + // for testing purposes, one portion is <= 2**28 elements + const PortionOffsetT PORTION_SIZE = ((1 << 28) - 1) / ONESWEEP_TILE_ITEMS * ONESWEEP_TILE_ITEMS; + int num_passes = ::cuda::ceil_div(end_bit - begin_bit, RADIX_BITS); + OffsetT num_portions = static_cast(::cuda::ceil_div(num_items, PORTION_SIZE)); + PortionOffsetT max_num_blocks = ::cuda::ceil_div( + static_cast(::cuda::std::min(num_items, static_cast(PORTION_SIZE))), ONESWEEP_TILE_ITEMS); + + size_t value_size = KEYS_ONLY ? 0 : kernel_source.ValueSize(); + size_t allocation_sizes[] = { + // bins + num_portions * num_passes * RADIX_DIGITS * sizeof(OffsetT), + // lookback + max_num_blocks * RADIX_DIGITS * sizeof(AtomicOffsetT), + // extra key buffer + is_overwrite_okay || num_passes <= 1 ? 0 : num_items * kernel_source.KeySize(), + // extra value buffer + is_overwrite_okay || num_passes <= 1 ? 0 : num_items * value_size, + // counters + num_portions * num_passes * sizeof(AtomicOffsetT), + }; + constexpr int NUM_ALLOCATIONS = sizeof(allocation_sizes) / sizeof(allocation_sizes[0]); + void* allocations[NUM_ALLOCATIONS] = {}; + if (const auto error = + detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)) + { + return error; + } + + // just return if no temporary storage is provided + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + OffsetT* d_bins = (OffsetT*) allocations[0]; + AtomicOffsetT* d_lookback = (AtomicOffsetT*) allocations[1]; + KeyT* d_keys_tmp2 = (KeyT*) allocations[2]; + ValueT* d_values_tmp2 = (ValueT*) allocations[3]; + AtomicOffsetT* d_ctrs = (AtomicOffsetT*) allocations[4]; + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + constexpr OffsetT pdl_max_items = static_cast(1) << 24; + const bool use_pdl = num_items <= pdl_max_items && cc >= ::cuda::compute_capability{9, 0}; + + const size_t num_counter_items = static_cast(num_portions) * num_passes; + const size_t num_bin_items = static_cast(num_passes) * RADIX_DIGITS; + int device = -1; + int num_sms = 0; + + if (const auto error = CubDebug(cudaGetDevice(&device))) + { + return error; + } + + if (const auto error = CubDebug(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device))) + { + return error; + } + + const int HISTO_BLOCK_THREADS = policy.histogram.threads_per_block; + int histo_blocks_per_sm = 1; + auto histogram_kernel = kernel_source.RadixSortHistogramKernel(); + + if (const auto error = + CubDebug(launcher_factory.MaxSmOccupancy(histo_blocks_per_sm, histogram_kernel, HISTO_BLOCK_THREADS, 0))) + { + return error; + } + +// log histogram_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking histogram_kernel<<<%d, %d, 0, %lld>>>(), %d items per iteration, " + "%d SM occupancy, bit_grain %d\n", + histo_blocks_per_sm * num_sms, + HISTO_BLOCK_THREADS, + reinterpret_cast(stream), + policy.histogram.items_per_thread, + histo_blocks_per_sm, + policy.histogram.radix_bits); +#endif + + // exclusive sums to determine starts + const int SCAN_BLOCK_THREADS = policy.exclusive_sum.threads_per_block; + +// log exclusive_sum_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking exclusive_sum_kernel<<<%d, %d, 0, %lld>>>(), bit_grain %d\n", + num_passes, + SCAN_BLOCK_THREADS, + reinterpret_cast(stream), + policy.exclusive_sum.radix_bits); +#endif + + // Initialization is intentionally adjacent to the histogram launch. For the PDL path, this avoids consuming the + // short init kernel's runtime in host-side launch setup work before the dependent histogram is submitted. + { + constexpr int init_startup_threads = 256; + const size_t num_init_items = ::cuda::std::max(num_counter_items, num_bin_items); + const int init_startup_blocks = + static_cast(::cuda::ceil_div(num_init_items, static_cast(init_startup_threads))); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_bins_and_counters_kernel<<<%d, %d, 0, %lld>>>()\n", + init_startup_blocks, + init_startup_threads, + reinterpret_cast(stream)); +#endif + + if (const auto error = CubDebug( + launcher_factory(init_startup_blocks, init_startup_threads, 0, stream, use_pdl) + .doit( + kernel_source.RadixSortInitBinsAndCountersKernel(), d_ctrs, num_counter_items, d_bins, num_bin_items))) + { + return error; + } + } + + if (const auto error = CubDebug( + launcher_factory(histo_blocks_per_sm * num_sms, HISTO_BLOCK_THREADS, 0, stream, use_pdl) + .doit(histogram_kernel, d_bins, d_keys.Current(), num_items, begin_bit, end_bit, decomposer))) + { + return error; + } + + if (const auto error = CubDebug(launcher_factory(num_passes, SCAN_BLOCK_THREADS, 0, stream, use_pdl) + .doit(kernel_source.RadixSortExclusiveSumKernel(), d_bins))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // use the other buffer if no overwrite is allowed + KeyT* d_keys_tmp = d_keys.Alternate(); + ValueT* d_values_tmp = d_values.Alternate(); + if (!is_overwrite_okay && num_passes % 2 == 0) + { + d_keys.d_buffers[1] = d_keys_tmp2; + d_values.d_buffers[1] = d_values_tmp2; + } + + for (int current_bit = begin_bit, pass = 0; current_bit < end_bit; current_bit += RADIX_BITS, ++pass) + { + int num_bits = ::cuda::std::min(end_bit - current_bit, RADIX_BITS); + for (OffsetT portion = 0; portion < num_portions; ++portion) + { + PortionOffsetT portion_num_items = static_cast( + ::cuda::std::min(num_items - portion * PORTION_SIZE, static_cast(PORTION_SIZE))); + + PortionOffsetT num_blocks = ::cuda::ceil_div(portion_num_items, ONESWEEP_TILE_ITEMS); + const size_t num_lookback_items = static_cast(num_blocks) * RADIX_DIGITS; + + if (use_pdl) + { + constexpr int init_lookback_threads = 256; + const int init_lookback_blocks = + static_cast(::cuda::ceil_div(num_lookback_items, static_cast(init_lookback_threads))); + + if (const auto error = CubDebug( + launcher_factory(init_lookback_blocks, init_lookback_threads, 0, stream, use_pdl) + .doit( + kernel_source.RadixSortInitLookbackKernel(), d_lookback, num_lookback_items, d_lookback, size_t{0}))) + { + return error; + } + } + else + { + if (const auto error = + CubDebug(cudaMemsetAsync(d_lookback, 0, num_lookback_items * sizeof(AtomicOffsetT), stream))) + { + return error; + } + } + +// log onesweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking onesweep_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, " + "current bit %d, bit_grain %d, portion %d/%d\n", + num_blocks, + ONESWEEP_BLOCK_THREADS, + reinterpret_cast(stream), + policy.onesweep.items_per_thread, + current_bit, + num_bits, + static_cast(portion), + static_cast(num_portions)); +#endif + + auto onesweep_kernel = kernel_source.RadixSortOnesweepKernel(); + + if (const auto error = CubDebug( + launcher_factory(num_blocks, ONESWEEP_BLOCK_THREADS, 0, stream, use_pdl) + .doit( + onesweep_kernel, + d_lookback, + d_ctrs + portion * num_passes + pass, + portion < num_portions - 1 ? d_bins + ((portion + 1) * num_passes + pass) * RADIX_DIGITS : nullptr, + d_bins + (portion * num_passes + pass) * RADIX_DIGITS, + d_keys.Alternate(), + kernel_source.AdvanceKeys(d_keys.Current(), portion * PORTION_SIZE), + d_values.Alternate(), + kernel_source.AdvanceValues(d_values.Current(), portion * PORTION_SIZE), + portion_num_items, + current_bit, + num_bits, + decomposer))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + // use the temporary buffers if no overwrite is allowed + if (!is_overwrite_okay && pass == 0) + { + d_keys = num_passes % 2 == 0 ? DoubleBuffer(d_keys_tmp, d_keys_tmp2) + : DoubleBuffer(d_keys_tmp2, d_keys_tmp); + d_values = num_passes % 2 == 0 ? DoubleBuffer(d_values_tmp, d_values_tmp2) + : DoubleBuffer(d_values_tmp2, d_values_tmp); + } + d_keys.selector ^= 1; + d_values.selector ^= 1; + } + + return cudaSuccess; + } + +public: + /** + * @brief Invocation (run multiple digit passes) + * + * @tparam ActivePolicyT + * Umbrella policy active for the target device + * + * @tparam UpsweepKernelT + * Function type of cub::DeviceRadixSortUpsweepKernel + * + * @tparam ScanKernelT + * Function type of cub::SpineScanKernel + * + * @tparam DownsweepKernelT + * Function type of cub::DeviceRadixSortDownsweepKernel + * + * @param[in] upsweep_kernel + * Kernel function pointer to parameterization of cub::DeviceRadixSortUpsweepKernel + * + * @param[in] alt_upsweep_kernel + * Alternate kernel function pointer to parameterization of cub::DeviceRadixSortUpsweepKernel + * + * @param[in] scan_kernel + * Kernel function pointer to parameterization of cub::SpineScanKernel + * + * @param[in] downsweep_kernel + * Kernel function pointer to parameterization of cub::DeviceRadixSortDownsweepKernel + * + * @param[in] alt_downsweep_kernel + * Alternate kernel function pointer to parameterization of + * cub::DeviceRadixSortDownsweepKernel + */ + // TODO(bgruber): Remove in CCCL 4.0 + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t InvokePasses( + UpsweepKernelT upsweep_kernel, + UpsweepKernelT alt_upsweep_kernel, + ScanKernelT scan_kernel, + DownsweepKernelT downsweep_kernel, + DownsweepKernelT alt_downsweep_kernel, + ActivePolicyT policy = {}) + { + return __invoke_passes( + upsweep_kernel, + alt_downsweep_kernel, + scan_kernel, + downsweep_kernel, + alt_downsweep_kernel, + detail::radix_sort::convert_policy(policy)); + } + +private: + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t __invoke_passes( + UpsweepKernelT upsweep_kernel, + UpsweepKernelT alt_upsweep_kernel, + ScanKernelT scan_kernel, + DownsweepKernelT downsweep_kernel, + DownsweepKernelT alt_downsweep_kernel, + const RadixSortPolicy& policy) + { + // Get device ordinal + int device_ordinal; + if (const auto error = CubDebug(cudaGetDevice(&device_ordinal))) + { + return error; + } + + // Get SM count + int sm_count; + if (const auto error = CubDebug(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_ordinal))) + { + return error; + } + + // Init regular and alternate-digit kernel configurations + PassConfig pass_config, alt_pass_config; + if (const auto error = pass_config.__init_pass_config( + upsweep_kernel, + scan_kernel, + downsweep_kernel, + sm_count, + num_items, + policy.downsweep.radix_bits, + policy.upsweep, + policy.scan, + policy.downsweep, + launcher_factory)) + { + return error; + } + + if (const auto error = alt_pass_config.__init_pass_config( + alt_upsweep_kernel, + scan_kernel, + alt_downsweep_kernel, + sm_count, + num_items, + policy.alt_downsweep.radix_bits, + policy.alt_upsweep, + policy.scan, + policy.alt_downsweep, + launcher_factory)) + { + return error; + } + + // Get maximum spine length + int max_grid_size = ::cuda::std::max(pass_config.max_downsweep_grid_size, alt_pass_config.max_downsweep_grid_size); + int spine_length = (max_grid_size * pass_config.radix_digits) + pass_config.scan_config.tile_size; + + // Temporary storage allocation requirements + void* allocations[3] = {}; + size_t allocation_sizes[3] = { + // bytes needed for privatized block digit histograms + spine_length * sizeof(OffsetT), + + // bytes needed for 3rd keys buffer + (is_overwrite_okay) ? 0 : num_items * kernel_source.KeySize(), + + // bytes needed for 3rd values buffer + (is_overwrite_okay || (KEYS_ONLY)) ? 0 : num_items * kernel_source.ValueSize(), + }; + + // Alias the temporary allocations from the single storage blob (or compute the necessary size of the blob) + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + // Pass planning. Run passes of the alternate digit-size configuration until we have an even multiple of our + // preferred digit size + int num_bits = end_bit - begin_bit; + int num_passes = ::cuda::ceil_div(num_bits, pass_config.radix_bits); + bool is_num_passes_odd = num_passes & 1; + int max_alt_passes = (num_passes * pass_config.radix_bits) - num_bits; + int alt_end_bit = ::cuda::std::min(end_bit, begin_bit + (max_alt_passes * alt_pass_config.radix_bits)); + + // Alias the temporary storage allocations + OffsetT* d_spine = static_cast(allocations[0]); + + DoubleBuffer d_keys_remaining_passes( + (is_overwrite_okay || is_num_passes_odd) ? d_keys.Alternate() : static_cast(allocations[1]), + (is_overwrite_okay) ? d_keys.Current() + : (is_num_passes_odd) ? static_cast(allocations[1]) + : d_keys.Alternate()); + + DoubleBuffer d_values_remaining_passes( + (is_overwrite_okay || is_num_passes_odd) ? d_values.Alternate() : static_cast(allocations[2]), + (is_overwrite_okay) ? d_values.Current() + : (is_num_passes_odd) ? static_cast(allocations[2]) + : d_values.Alternate()); + + // Run first pass, consuming from the input's current buffers + int current_bit = begin_bit; + if (const auto error = CubDebug(InvokePass( + d_keys.Current(), + d_keys_remaining_passes.Current(), + d_values.Current(), + d_values_remaining_passes.Current(), + d_spine, + spine_length, + current_bit, + (current_bit < alt_end_bit) ? alt_pass_config : pass_config))) + { + return error; + } + + // Run remaining passes + while (current_bit < end_bit) + { + if (const auto error = CubDebug(InvokePass( + d_keys_remaining_passes.d_buffers[d_keys_remaining_passes.selector], + d_keys_remaining_passes.d_buffers[d_keys_remaining_passes.selector ^ 1], + d_values_remaining_passes.d_buffers[d_keys_remaining_passes.selector], + d_values_remaining_passes.d_buffers[d_keys_remaining_passes.selector ^ 1], + d_spine, + spine_length, + current_bit, + (current_bit < alt_end_bit) ? alt_pass_config : pass_config))) + { + return error; + } + + // Invert selectors + d_keys_remaining_passes.selector ^= 1; + d_values_remaining_passes.selector ^= 1; + } + + // Update selector + if (!is_overwrite_okay) + { + num_passes = 1; // Sorted data always ends up in the other vector + } + + d_keys.selector = (d_keys.selector + num_passes) & 1; + d_values.selector = (d_values.selector + num_passes) & 1; + + return cudaSuccess; + } + +public: + // TODO(bgruber): Remove in CCCL 4.0 + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t InvokeCopy() + { + // is_overwrite_okay == false here + // Return the number of temporary bytes if requested + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + +// Copy keys +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking async copy of %lld keys on stream %lld\n", (long long) num_items, (long long) stream); +#endif + if (const auto error = CubDebug(cudaMemcpyAsync( + d_keys.Alternate(), d_keys.Current(), num_items * kernel_source.KeySize(), cudaMemcpyDefault, stream))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + d_keys.selector ^= 1; + + // Copy values if necessary + if constexpr (!KEYS_ONLY) + { +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking async copy of %lld values on stream %lld\n", (long long) num_items, (long long) stream); +#endif + if (const auto error = CubDebug(cudaMemcpyAsync( + d_values.Alternate(), d_values.Current(), num_items * kernel_source.ValueSize(), cudaMemcpyDefault, stream))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + d_values.selector ^= 1; + + return cudaSuccess; + } + + // TODO(bgruber): Remove in CCCL 4.0 + /// Invocation + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT = {}) + { + struct policy_getter + { + _CCCL_HOST_DEVICE_API _CCCL_FORCEINLINE constexpr auto operator()() const + { + return detail::radix_sort::convert_policy(); + } + }; + return __invoke(policy_getter{}); + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t __invoke(PolicyGetter policy_getter) + { + CUB_DETAIL_CONSTEXPR_ISH auto policy = policy_getter(); + + // Return if empty problem, or if no bits to sort and double-buffering is used + if (num_items == 0 || (begin_bit == end_bit && is_overwrite_okay)) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + // Check if simple copy suffices (is_overwrite_okay == false at this point) + if (begin_bit == end_bit) + { + bool has_uva = false; + if (const auto error = detail::HasUVA(has_uva)) + { + return error; + } + if (has_uva) + { + return InvokeCopy(); + } + } + + // Older GCCs choke on using policy in a constexpr expression even if it is constexpr itself +#if defined(CUB_DEFINE_RUNTIME_POLICIES) || _CCCL_COMPILER(GCC, <, 10) + const auto tile_items = + static_cast(policy.single_tile.threads_per_block * policy.single_tile.items_per_thread); +#else // ^^^ runtime tile items ^^^ / vvv constexpr tile items vvv + constexpr auto tile_items = OffsetT{policy.single_tile.threads_per_block * policy.single_tile.items_per_thread}; +#endif // ^^^ true constexpr tile_items ^^^ + + // Force kernel code-generation in all compiler passes + if (num_items <= tile_items) + { + // Small, single tile size + return __invoke_single_tile(kernel_source.RadixSortSingleTileKernel(), policy.single_tile); + } + +#if _CCCL_COMPILER(GCC, <, 10) + // gcc 7-9 fail to use `policy` in a constant expression, so we just compute it again inplace + if CUB_DETAIL_CONSTEXPR_ISH (PolicyGetter{}().algorithm == RadixSortAlgorithm::onesweep) +#else // _CCCL_COMPILER(GCC, <, 8) + if CUB_DETAIL_CONSTEXPR_ISH (policy.algorithm == RadixSortAlgorithm::onesweep) +#endif // _CCCL_COMPILER(GCC, <, 8) + { + return __invoke_onesweep(policy); + } + else + { + return __invoke_passes( + kernel_source.RadixSortUpsweepKernel(), + kernel_source.RadixSortAltUpsweepKernel(), + kernel_source.DeviceRadixSortScanBinsKernel(), + kernel_source.RadixSortDownsweepKernel(), + kernel_source.RadixSortAltDownsweepKernel(), + policy); + } + } + + //------------------------------------------------------------------------------ + // Dispatch entrypoints + //------------------------------------------------------------------------------ + + /** + * @brief Internal dispatch routine + * + * @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,out] d_keys + * Double-buffer whose current buffer contains the unsorted input keys and, + * upon return, is updated to point to the sorted output keys + * + * @param[in,out] d_values + * Double-buffer whose current buffer contains the unsorted input values and, + * upon return, is updated to point to the sorted output values + * + * @param[in] num_items + * Number of items to sort + * + * @param[in] begin_bit + * The beginning (least-significant) bit index needed for key comparison + * + * @param[in] end_bit + * The past-the-end (most-significant) bit index needed for key comparison + * + * @param[in] is_overwrite_okay + * Whether is okay to overwrite source buffers + * + * @param[in] stream + * CUDA stream to launch kernels within. Default is stream0. + */ + // TODO(bgruber): Remove in CCCL 4.0 + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + OffsetT num_items, + int begin_bit, + int end_bit, + bool is_overwrite_okay, + cudaStream_t stream, + DecomposerT decomposer = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) + { + // Get PTX version + int ptx_version = 0; + if (const auto error = CubDebug(launcher_factory.PtxVersion(ptx_version))) + { + return error; + } + + // Create dispatch functor + DispatchRadixSort dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + begin_bit, + end_bit, + is_overwrite_okay, + stream, + ptx_version, + decomposer, + kernel_source, + launcher_factory); + + // Dispatch to chained policy + if (const auto error = CubDebug(max_policy.Invoke(ptx_version, dispatch))) + { + return error; + } + + return cudaSuccess; + } +}; + +namespace detail::radix_sort +{ +template +struct pass_config +{ + UpsweepKernelT upsweep_kernel; + KernelConfig upsweep_config; + ScanKernelT scan_kernel; + KernelConfig scan_config; + DownsweepKernelT downsweep_kernel; + KernelConfig downsweep_config; + int radix_bits; + int radix_digits; + int max_downsweep_grid_size; + GridEvenShare even_share; + + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t init( + UpsweepKernelT upsweep_kern, + ScanKernelT scan_kern, + DownsweepKernelT downsweep_kern, + int sm_count, + OffsetT num_items, + int pass_radix_bits, + RadixSortUpsweepPolicy upsweep_policy, + ScanPolicy scan_policy, + RadixSortDownsweepPolicy downsweep_policy, + KernelLauncherFactory launcher_factory) + { + this->upsweep_kernel = upsweep_kern; + this->scan_kernel = scan_kern; + this->downsweep_kernel = downsweep_kern; + this->radix_bits = pass_radix_bits; + radix_digits = 1 << radix_bits; + + if (const auto error = CubDebug(upsweep_config.__init(upsweep_kernel, upsweep_policy, launcher_factory))) + { + return error; + } + + if (const auto error = CubDebug(scan_config.__init(scan_kernel, scan_policy.lookback, launcher_factory))) + { + return error; + } + + if (const auto error = CubDebug(downsweep_config.__init(downsweep_kernel, downsweep_policy, launcher_factory))) + { + return error; + } + + max_downsweep_grid_size = (downsweep_config.sm_occupancy * sm_count) * detail::subscription_factor; + + even_share.DispatchInit( + num_items, max_downsweep_grid_size, ::cuda::std::max(downsweep_config.tile_size, upsweep_config.tile_size)); + + return cudaSuccess; + } +}; + +template +struct dispatch_impl +{ + static constexpr bool keys_only = ::cuda::std::is_same_v; + + void* d_temp_storage; + size_t& temp_storage_bytes; + DoubleBuffer& d_keys; // current buffer holds unsorted input keys, updated before returning to caller + DoubleBuffer& d_values; // current buffer holds unsorted input keys, updated before returning to caller + OffsetT num_items; + int begin_bit; // the beginning least-significant bit index for key comparison + int end_bit; // the one-past-the-end least-significant bit index for key comparison + bool can_overwrite_source_buffer; + cudaStream_t stream; + DecomposerT decomposer; + KernelSource kernel_source; + KernelLauncherFactory launcher_factory; + + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t + invoke_single_tile(SingleTileKernelT single_tile_kernel, RadixSortDownsweepPolicy policy) + { + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + // Log single_tile_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking single_tile_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, %d SM occupancy, current bit " + "%d, bit_grain %d\n", + 1, + policy.threads_per_block, + (long long) stream, + policy.items_per_thread, + 1, + begin_bit, + policy.radix_bits); +#endif + + // Invoke upsweep_kernel with same grid size as downsweep_kernel + if (const auto error = CubDebug( + launcher_factory(1, policy.threads_per_block, 0, stream) + .doit(single_tile_kernel, + d_keys.Current(), + d_keys.Alternate(), + d_values.Current(), + d_values.Alternate(), + num_items, + begin_bit, + end_bit, + decomposer))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Update selector + d_keys.selector ^= 1; + d_values.selector ^= 1; + + return cudaSuccess; + } + + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_copy() + { + // can_overwrite_source_buffer == false here + // Return the number of temporary bytes if requested + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + +// Copy keys +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking async copy of %lld keys on stream %lld\n", (long long) num_items, (long long) stream); +#endif + if (const auto error = CubDebug(cudaMemcpyAsync( + d_keys.Alternate(), d_keys.Current(), num_items * kernel_source.KeySize(), cudaMemcpyDefault, stream))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + d_keys.selector ^= 1; + + // Copy values if necessary + if constexpr (!keys_only) + { +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking async copy of %lld values on stream %lld\n", (long long) num_items, (long long) stream); +#endif + if (const auto error = CubDebug(cudaMemcpyAsync( + d_values.Alternate(), d_values.Current(), num_items * kernel_source.ValueSize(), cudaMemcpyDefault, stream))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + d_values.selector ^= 1; + + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_pass( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + OffsetT* d_spine, + int /*spine_length*/, + int& current_bit, + PassConfigT& pass_config) + { + int pass_bits = ::cuda::std::min(pass_config.radix_bits, end_bit - current_bit); + +// Log upsweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking upsweep_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, %d SM occupancy, current bit %d, " + "bit_grain %d\n", + pass_config.even_share.grid_size, + pass_config.upsweep_config.threads_per_block, + (long long) stream, + pass_config.upsweep_config.items_per_thread, + pass_config.upsweep_config.sm_occupancy, + current_bit, + pass_bits); +#endif + + // Spine length written by the upsweep kernel in the current pass. + int pass_spine_length = pass_config.even_share.grid_size * pass_config.radix_digits; + + // Invoke upsweep_kernel with same grid size as downsweep_kernel + if (const auto error = CubDebug( + launcher_factory(pass_config.even_share.grid_size, pass_config.upsweep_config.threads_per_block, 0, stream) + .doit(pass_config.upsweep_kernel, + d_keys_in, + d_spine, + num_items, + current_bit, + pass_bits, + pass_config.even_share, + decomposer))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + +// Log scan_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking scan_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread\n", + 1, + pass_config.scan_config.threads_per_block, + (long long) stream, + pass_config.scan_config.items_per_thread); +#endif + + // Invoke scan_kernel + if (const auto error = CubDebug(launcher_factory(1, pass_config.scan_config.threads_per_block, 0, stream) + .doit(pass_config.scan_kernel, d_spine, pass_spine_length))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + +// Log downsweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking downsweep_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, %d SM occupancy\n", + pass_config.even_share.grid_size, + pass_config.downsweep_config.threads_per_block, + (long long) stream, + pass_config.downsweep_config.items_per_thread, + pass_config.downsweep_config.sm_occupancy); +#endif + + // Invoke downsweep_kernel + if (const auto error = CubDebug( + launcher_factory(pass_config.even_share.grid_size, pass_config.downsweep_config.threads_per_block, 0, stream) + .doit(pass_config.downsweep_kernel, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + d_spine, + num_items, + current_bit, + pass_bits, + pass_config.even_share, + decomposer))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Update current bit + current_bit += pass_bits; + + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invoke_passes( + UpsweepKernelT upsweep_kernel, + UpsweepKernelT alt_upsweep_kernel, + ScanKernelT scan_kernel, + DownsweepKernelT downsweep_kernel, + DownsweepKernelT alt_downsweep_kernel, + const RadixSortPolicy& policy) + { + // Get device ordinal + int device_ordinal; + if (const auto error = CubDebug(cudaGetDevice(&device_ordinal))) + { + return error; + } + + // Get SM count + int sm_count; + if (const auto error = CubDebug(cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, device_ordinal))) + { + return error; + } + + // Init regular and alternate-digit kernel configurations + pass_config pc, alt_pc; + if (const auto error = pc.init( + upsweep_kernel, + scan_kernel, + downsweep_kernel, + sm_count, + num_items, + policy.downsweep.radix_bits, + policy.upsweep, + policy.scan, + policy.downsweep, + launcher_factory)) + { + return error; + } + + if (const auto error = alt_pc.init( + alt_upsweep_kernel, + scan_kernel, + alt_downsweep_kernel, + sm_count, + num_items, + policy.alt_downsweep.radix_bits, + policy.alt_upsweep, + policy.scan, + policy.alt_downsweep, + launcher_factory)) + { + return error; + } + + // Get maximum spine length + int max_grid_size = ::cuda::std::max(pc.max_downsweep_grid_size, alt_pc.max_downsweep_grid_size); + int spine_length = (max_grid_size * pc.radix_digits) + pc.scan_config.tile_size; + + // Temporary storage allocation requirements + void* allocations[3] = {}; + size_t allocation_sizes[3] = { + // bytes needed for privatized block digit histograms + spine_length * sizeof(OffsetT), + + // bytes needed for 3rd keys buffer + (can_overwrite_source_buffer) ? 0 : num_items * kernel_source.KeySize(), + + // bytes needed for 3rd values buffer + (can_overwrite_source_buffer || (keys_only)) ? 0 : num_items * kernel_source.ValueSize(), + }; + + // Alias the temporary allocations from the single storage blob (or compute the necessary size of the blob) + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + // Pass planning. Run passes of the alternate digit-size configuration until we have an even multiple of our + // preferred digit size + int num_bits = end_bit - begin_bit; + int num_passes = ::cuda::ceil_div(num_bits, pc.radix_bits); + bool is_num_passes_odd = num_passes & 1; + int max_alt_passes = (num_passes * pc.radix_bits) - num_bits; + int alt_end_bit = ::cuda::std::min(end_bit, begin_bit + (max_alt_passes * alt_pc.radix_bits)); + + // Alias the temporary storage allocations + OffsetT* d_spine = static_cast(allocations[0]); + + DoubleBuffer d_keys_remaining_passes( + (can_overwrite_source_buffer || is_num_passes_odd) ? d_keys.Alternate() : static_cast(allocations[1]), + (can_overwrite_source_buffer) ? d_keys.Current() + : (is_num_passes_odd) ? static_cast(allocations[1]) + : d_keys.Alternate()); + + DoubleBuffer d_values_remaining_passes( + (can_overwrite_source_buffer || is_num_passes_odd) ? d_values.Alternate() : static_cast(allocations[2]), + (can_overwrite_source_buffer) ? d_values.Current() + : (is_num_passes_odd) ? static_cast(allocations[2]) + : d_values.Alternate()); + + // Run first pass, consuming from the input's current buffers + int current_bit = begin_bit; + if (const auto error = CubDebug(invoke_pass( + d_keys.Current(), + d_keys_remaining_passes.Current(), + d_values.Current(), + d_values_remaining_passes.Current(), + d_spine, + spine_length, + current_bit, + (current_bit < alt_end_bit) ? alt_pc : pc))) + { + return error; + } + + // Run remaining passes + while (current_bit < end_bit) + { + if (const auto error = CubDebug(invoke_pass( + d_keys_remaining_passes.d_buffers[d_keys_remaining_passes.selector], + d_keys_remaining_passes.d_buffers[d_keys_remaining_passes.selector ^ 1], + d_values_remaining_passes.d_buffers[d_keys_remaining_passes.selector], + d_values_remaining_passes.d_buffers[d_keys_remaining_passes.selector ^ 1], + d_spine, + spine_length, + current_bit, + (current_bit < alt_end_bit) ? alt_pc : pc))) + { + return error; + } + + // Invert selectors + d_keys_remaining_passes.selector ^= 1; + d_values_remaining_passes.selector ^= 1; + } + + // Update selector + if (!can_overwrite_source_buffer) + { + num_passes = 1; // Sorted data always ends up in the other vector + } + + d_keys.selector = (d_keys.selector + num_passes) & 1; + d_values.selector = (d_values.selector + num_passes) & 1; + + return cudaSuccess; + } + + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_onesweep(RadixSortPolicy policy) + { + // PortionOffsetT is used for offsets within a portion, and must be signed. + using PortionOffsetT = int; + using AtomicOffsetT = PortionOffsetT; + + // compute temporary storage size + const int radix_bits = policy.onesweep.radix_bits; + const int radix_digits = 1 << radix_bits; + const int onesweep_items_per_thread = policy.onesweep.items_per_thread; + const int onesweep_block_threads = policy.onesweep.threads_per_block; + const int onesweep_tile_items = onesweep_items_per_thread * onesweep_block_threads; + // portions handle inputs with >=2**30 elements, due to the way lookback works + // for testing purposes, one portion is <= 2**28 elements + const PortionOffsetT portion_size = ((1 << 28) - 1) / onesweep_tile_items * onesweep_tile_items; + int num_passes = ::cuda::ceil_div(end_bit - begin_bit, radix_bits); + OffsetT num_portions = static_cast(::cuda::ceil_div(num_items, portion_size)); + PortionOffsetT max_num_blocks = ::cuda::ceil_div( + static_cast(::cuda::std::min(num_items, static_cast(portion_size))), onesweep_tile_items); + + size_t value_size = keys_only ? 0 : kernel_source.ValueSize(); + size_t allocation_sizes[] = { + // bins + num_portions * num_passes * radix_digits * sizeof(OffsetT), + // lookback + max_num_blocks * radix_digits * sizeof(AtomicOffsetT), + // extra key buffer + can_overwrite_source_buffer || num_passes <= 1 ? 0 : num_items * kernel_source.KeySize(), + // extra value buffer + can_overwrite_source_buffer || num_passes <= 1 ? 0 : num_items * value_size, + // counters + num_portions * num_passes * sizeof(AtomicOffsetT), + }; + constexpr int num_allocations = sizeof(allocation_sizes) / sizeof(allocation_sizes[0]); + void* allocations[num_allocations] = {}; + if (const auto error = + detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)) + { + return error; + } + + // just return if no temporary storage is provided + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + OffsetT* d_bins = (OffsetT*) allocations[0]; + AtomicOffsetT* d_lookback = (AtomicOffsetT*) allocations[1]; + KeyT* d_keys_tmp2 = (KeyT*) allocations[2]; + ValueT* d_values_tmp2 = (ValueT*) allocations[3]; + AtomicOffsetT* d_ctrs = (AtomicOffsetT*) allocations[4]; + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + constexpr OffsetT pdl_max_items = static_cast(1) << 24; + const bool use_pdl = num_items <= pdl_max_items && cc >= ::cuda::compute_capability{9, 0}; + + const size_t num_counter_items = static_cast(num_portions) * num_passes; + const size_t num_bin_items = static_cast(num_passes) * radix_digits; + int device = -1; + int num_sms = 0; + + if (const auto error = CubDebug(cudaGetDevice(&device))) + { + return error; + } + + if (const auto error = CubDebug(cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, device))) + { + return error; + } + + const int histo_block_threads = policy.histogram.threads_per_block; + int histo_blocks_per_sm = 1; + auto histogram_kernel = kernel_source.RadixSortHistogramKernel(); + + if (const auto error = + CubDebug(launcher_factory.MaxSmOccupancy(histo_blocks_per_sm, histogram_kernel, histo_block_threads, 0))) + { + return error; + } + +// log histogram_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking histogram_kernel<<<%d, %d, 0, %lld>>>(), %d items per iteration, " + "%d SM occupancy, bit_grain %d\n", + histo_blocks_per_sm * num_sms, + histo_block_threads, + reinterpret_cast(stream), + policy.histogram.items_per_thread, + histo_blocks_per_sm, + policy.histogram.radix_bits); +#endif + + // exclusive sums to determine starts + const int scan_block_threads = policy.exclusive_sum.threads_per_block; + +// log exclusive_sum_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking exclusive_sum_kernel<<<%d, %d, 0, %lld>>>(), bit_grain %d\n", + num_passes, + scan_block_threads, + reinterpret_cast(stream), + policy.exclusive_sum.radix_bits); +#endif + + // Initialization is intentionally adjacent to the histogram launch. For the PDL path, this avoids consuming the + // short init kernel's runtime in host-side launch setup work before the dependent histogram is submitted. + { + constexpr int init_startup_threads = 256; + const size_t num_init_items = ::cuda::std::max(num_counter_items, num_bin_items); + const int init_startup_blocks = + static_cast(::cuda::ceil_div(num_init_items, static_cast(init_startup_threads))); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_bins_and_counters_kernel<<<%d, %d, 0, %lld>>>()\n", + init_startup_blocks, + init_startup_threads, + reinterpret_cast(stream)); +#endif + + if (const auto error = CubDebug( + launcher_factory(init_startup_blocks, init_startup_threads, 0, stream, use_pdl) + .doit( + kernel_source.RadixSortInitBinsAndCountersKernel(), d_ctrs, num_counter_items, d_bins, num_bin_items))) + { + return error; + } + } + + if (const auto error = CubDebug( + launcher_factory(histo_blocks_per_sm * num_sms, histo_block_threads, 0, stream, use_pdl) + .doit(histogram_kernel, d_bins, d_keys.Current(), num_items, begin_bit, end_bit, decomposer))) + { + return error; + } + + if (const auto error = CubDebug(launcher_factory(num_passes, scan_block_threads, 0, stream, use_pdl) + .doit(kernel_source.RadixSortExclusiveSumKernel(), d_bins))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // use the other buffer if no overwrite is allowed + KeyT* d_keys_tmp = d_keys.Alternate(); + ValueT* d_values_tmp = d_values.Alternate(); + if (!can_overwrite_source_buffer && num_passes % 2 == 0) + { + d_keys.d_buffers[1] = d_keys_tmp2; + d_values.d_buffers[1] = d_values_tmp2; + } + + for (int current_bit = begin_bit, pass = 0; current_bit < end_bit; current_bit += radix_bits, ++pass) + { + int num_bits = ::cuda::std::min(end_bit - current_bit, radix_bits); + for (OffsetT portion = 0; portion < num_portions; ++portion) + { + PortionOffsetT portion_num_items = static_cast( + ::cuda::std::min(num_items - portion * portion_size, static_cast(portion_size))); + + PortionOffsetT num_blocks = ::cuda::ceil_div(portion_num_items, onesweep_tile_items); + const size_t num_lookback_items = static_cast(num_blocks) * radix_digits; + + if (use_pdl) + { + constexpr int init_lookback_threads = 256; + const int init_lookback_blocks = + static_cast(::cuda::ceil_div(num_lookback_items, static_cast(init_lookback_threads))); + + if (const auto error = CubDebug( + launcher_factory(init_lookback_blocks, init_lookback_threads, 0, stream, use_pdl) + .doit( + kernel_source.RadixSortInitLookbackKernel(), d_lookback, num_lookback_items, d_lookback, size_t{0}))) + { + return error; + } + } + else + { + if (const auto error = + CubDebug(cudaMemsetAsync(d_lookback, 0, num_lookback_items * sizeof(AtomicOffsetT), stream))) + { + return error; + } + } + +// log onesweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking onesweep_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, " + "current bit %d, bit_grain %d, portion %d/%d\n", + num_blocks, + onesweep_block_threads, + reinterpret_cast(stream), + policy.onesweep.items_per_thread, + current_bit, + num_bits, + static_cast(portion), + static_cast(num_portions)); +#endif + + auto onesweep_kernel = kernel_source.RadixSortOnesweepKernel(); + + if (const auto error = CubDebug( + launcher_factory(num_blocks, onesweep_block_threads, 0, stream, use_pdl) + .doit( + onesweep_kernel, + d_lookback, + d_ctrs + portion * num_passes + pass, + portion < num_portions - 1 ? d_bins + ((portion + 1) * num_passes + pass) * radix_digits : nullptr, + d_bins + (portion * num_passes + pass) * radix_digits, + d_keys.Alternate(), + kernel_source.AdvanceKeys(d_keys.Current(), portion * portion_size), + d_values.Alternate(), + kernel_source.AdvanceValues(d_values.Current(), portion * portion_size), + portion_num_items, + current_bit, + num_bits, + decomposer))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + // use the temporary buffers if no overwrite is allowed + if (!can_overwrite_source_buffer && pass == 0) + { + d_keys = num_passes % 2 == 0 ? DoubleBuffer(d_keys_tmp, d_keys_tmp2) + : DoubleBuffer(d_keys_tmp2, d_keys_tmp); + d_values = num_passes % 2 == 0 ? DoubleBuffer(d_values_tmp, d_values_tmp2) + : DoubleBuffer(d_values_tmp2, d_values_tmp); + } + d_keys.selector ^= 1; + d_values.selector ^= 1; + } + + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke(PolicyGetter policy_getter) + { + CUB_DETAIL_CONSTEXPR_ISH auto policy = policy_getter(); + + // Return if empty problem, or if no bits to sort and double-buffering is used + if (num_items == 0 || (begin_bit == end_bit && can_overwrite_source_buffer)) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + // Check if simple copy suffices (can_overwrite_source_buffer == false at this point) + if (begin_bit == end_bit) + { + bool has_uva = false; + if (const auto error = detail::HasUVA(has_uva)) + { + return error; + } + if (has_uva) + { + return invoke_copy(); + } + } + + // Force kernel code-generation in all compiler passes + if (num_items <= (static_cast(policy.single_tile.threads_per_block) * policy.single_tile.items_per_thread)) + { + // Small, single tile size + return invoke_single_tile(kernel_source.RadixSortSingleTileKernel(), policy.single_tile); + } + +#if _CCCL_COMPILER(GCC, <, 10) + // gcc 7-9 fail to use `policy` in a constant expression, so we just compute it again inplace + if CUB_DETAIL_CONSTEXPR_ISH (PolicyGetter{}().algorithm == RadixSortAlgorithm::onesweep) +#else // _CCCL_COMPILER(GCC, <, 10) + if CUB_DETAIL_CONSTEXPR_ISH (policy.algorithm == RadixSortAlgorithm::onesweep) +#endif // _CCCL_COMPILER(GCC, <, 10) + { + return invoke_onesweep(policy); + } + else + { + return invoke_passes( + kernel_source.RadixSortUpsweepKernel(), + kernel_source.RadixSortAltUpsweepKernel(), + kernel_source.DeviceRadixSortScanBinsKernel(), + kernel_source.RadixSortDownsweepKernel(), + kernel_source.RadixSortAltDownsweepKernel(), + policy); + } + } +}; + +template , + typename KernelSource = DeviceRadixSortKernelSource, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + OffsetT num_items, + int begin_bit, + int end_bit, + bool can_overwrite_source_buffer, + cudaStream_t stream, + DecomposerT decomposer = {}, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << policy_selector(cc); + _CubLog("Dispatching DeviceRadixSort to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + dispatch_impl impl{ + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + begin_bit, + end_bit, + can_overwrite_source_buffer, + stream, + decomposer, + kernel_source, + launcher_factory}; + + return dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) { + return impl.invoke(policy_getter); + }); +} +} // namespace detail::radix_sort + +CUB_NAMESPACE_END + +_CCCL_DIAG_POP diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce.cuh new file mode 100644 index 00000000..d8a11df4 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce.cuh @@ -0,0 +1,1057 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * @brief cub::DeviceReduce provides device-wide, parallel operations for + * computing a reduction across a sequence of data items residing within + * device-accessible memory. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include // for cub::detail::non_void_value_t, cub::detail::it_value_t + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// TODO(bgruber): included to not break users when moving DeviceSegmentedReduce to its own file. Remove in CCCL 4.0. +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::reduce +{ +template +struct DeviceReduceKernelSource +{ + // PolicySelector must be stateless, so we can pass the type to the kernel + static_assert(::cuda::std::is_empty_v); + + CUB_DEFINE_KERNEL_GETTER( + SingleTileKernel, + DeviceReduceSingleTileKernel) + + // The atomic code path finishes in one kernel, the two-phase code path writes to an intermediate buffer of + // accumulators + using reduce_kernel_output_t = ::cuda::std::conditional_t; + CUB_DEFINE_KERNEL_GETTER( + ReductionKernel, + DeviceReduceKernel) + + CUB_DEFINE_KERNEL_GETTER( + SingleTileSecondKernel, + DeviceReduceSingleTileKernel) + + CUB_DEFINE_KERNEL_GETTER( + DeferredSingleTileSecondKernel, + DeviceReduceDeferredSingleTileKernel< + PolicySelector, + AccumT*, + OutputIteratorT, + OffsetT, + KernelNumItemsT, + ReductionOpT, + InitValueT, + AccumT>) + + CUB_RUNTIME_FUNCTION static constexpr size_t AccumSize() + { + return sizeof(AccumT); + } + + CUB_RUNTIME_FUNCTION static constexpr size_t InitSize() + { + return sizeof(InitValueT); + } +}; + +// TODO(bgruber): remove in CCCL 4.0 +template +struct policy_selector_from_hub +{ + // this is only called in device code, so we can ignore the arch parameter + _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> ReducePolicy + { + using ap = typename PolicyHub::MaxPolicy::ActivePolicy; + using ap_reduce = typename ap::ReducePolicy; + using ap_single_tile = typename ap::SingleTilePolicy; + return ReducePolicy{ + ReducePassPolicy{ + ap_reduce::BLOCK_THREADS, + ap_reduce::ITEMS_PER_THREAD, + ap_reduce::VECTOR_LOAD_LENGTH, + ap_reduce::BLOCK_ALGORITHM, + ap_reduce::LOAD_MODIFIER, + }, + ReducePassPolicy{ + ap_single_tile::BLOCK_THREADS, + ap_single_tile::ITEMS_PER_THREAD, + ap_single_tile::VECTOR_LOAD_LENGTH, + ap_single_tile::BLOCK_ALGORITHM, + ap_single_tile::LOAD_MODIFIER, + }}; + } +}; +} // namespace detail::reduce + +/****************************************************************************** + * Single-problem dispatch + *****************************************************************************/ + +// TODO(bgruber): drop in CCCL 4.0 +/** + * @brief Utility class for dispatching the appropriately-tuned kernels for + * device-wide reduction + * + * Deprecated [Since 3.5] + * + * @tparam InputIteratorT + * Random-access input iterator type for reading input items @iterator + * + * @tparam OutputIteratorT + * Output iterator type for recording the reduced aggregate @iterator + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam ReductionOpT + * Binary reduction functor type having member + * `auto operator()(const T &a, const U &b)` + * + * @tparam InitValueT + * Initial value type + */ +template < + typename InputIteratorT, + typename OutputIteratorT, + typename OffsetT, + typename ReductionOpT, + typename InitValueT = cub::detail::non_void_value_t>, + typename AccumT = ::cuda::std::__accumulator_t, InitValueT>, + typename TransformOpT = ::cuda::std::identity, + typename PolicyHub = detail::reduce::policy_hub, + typename KernelSource = detail::reduce::DeviceReduceKernelSource< + detail::reduce::policy_selector_from_hub, + InputIteratorT, + OutputIteratorT, + OffsetT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT, + TransformOpT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce") DispatchReduce +{ + //--------------------------------------------------------------------------- + // Problem state + //--------------------------------------------------------------------------- + + /// Device-accessible allocation of temporary storage. When `nullptr`, the + /// required allocation size is written to `temp_storage_bytes` and no work + /// is done. + void* d_temp_storage; + + /// Reference to size in bytes of `d_temp_storage` allocation + size_t& temp_storage_bytes; + + /// Pointer to the input sequence of data items + InputIteratorT d_in; + + /// Pointer to the output aggregate + OutputIteratorT d_out; + + /// Total number of input items (i.e., length of `d_in`) + OffsetT num_items; + + /// Binary reduction functor + ReductionOpT reduction_op; + + /// The initial value of the reduction + InitValueT init; + + /// CUDA stream to launch kernels within. Default is stream0. + cudaStream_t stream; + + int ptx_version; + + TransformOpT transform_op; + + KernelSource kernel_source; + + KernelLauncherFactory launcher_factory; + + //--------------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------------- + + /// Constructor + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchReduce( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + int ptx_version, + TransformOpT transform_op = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_in(d_in) + , d_out(d_out) + , num_items(num_items) + , reduction_op(reduction_op) + , init(init) + , stream(stream) + , ptx_version(ptx_version) + , transform_op(transform_op) + , kernel_source(kernel_source) + , launcher_factory(launcher_factory) + {} + + //--------------------------------------------------------------------------- + // Small-problem (single tile) invocation + //--------------------------------------------------------------------------- + + /** + * @brief Invoke a single block block to reduce in-core + * + * @tparam ActivePolicyT + * Umbrella policy active for the target device + * + * @tparam SingleTileKernelT + * Function type of cub::DeviceReduceSingleTileKernel + * + * @param[in] single_tile_kernel + * Kernel function pointer to parameterization of + * cub::DeviceReduceSingleTileKernel + */ + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t + InvokeSingleTile(SingleTileKernelT single_tile_kernel, ActivePolicyT policy = {}) + { + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + +// Log single_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceReduceSingleTileKernel<<<1, %d, 0, %lld>>>(), " + "%d items per thread\n", + policy.SingleTile().ThreadsPerBlock(), + (long long) stream, + policy.SingleTile().ItemsPerThread()); +#endif // CUB_DEBUG_LOG + + // Invoke single_reduce_sweep_kernel + launcher_factory(1, policy.SingleTile().ThreadsPerBlock(), 0, stream) + .doit(single_tile_kernel, d_in, d_out, num_items, reduction_op, init, transform_op); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + return CubDebug(detail::DebugSyncStream(stream)); + } + + //--------------------------------------------------------------------------- + // Normal problem size invocation (two-pass) + //--------------------------------------------------------------------------- + + /** + * @brief Invoke two-passes to reduce + * @tparam ActivePolicyT + * Umbrella policy active for the target device + * + * @tparam ReduceKernelT + * Function type of cub::DeviceReduceKernel + * + * @tparam SingleTileKernelT + * Function type of cub::DeviceReduceSingleTileKernel + * + * @param[in] reduce_kernel + * Kernel function pointer to parameterization of cub::DeviceReduceKernel + * + * @param[in] single_tile_kernel + * Kernel function pointer to parameterization of + * cub::DeviceReduceSingleTileKernel + */ + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t + InvokePasses(ReduceKernelT reduce_kernel, SingleTileKernelT single_tile_kernel, ActivePolicyT active_policy = {}) + { + // Get SM count + int sm_count; + if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) + { + return error; + } + + // Init regular kernel configuration + detail::KernelConfig reduce_config; + if (const auto error = CubDebug(reduce_config.Init(reduce_kernel, active_policy.Reduce(), launcher_factory))) + { + return error; + } + + int reduce_device_occupancy = reduce_config.sm_occupancy * sm_count; + + // Even-share work distribution + int max_blocks = reduce_device_occupancy * detail::subscription_factor; + GridEvenShare even_share; + even_share.DispatchInit(num_items, max_blocks, reduce_config.tile_size); + + // Temporary storage allocation requirements + void* allocations[1] = {}; + size_t allocation_sizes[1] = { + max_blocks * kernel_source.AccumSize() // bytes needed for privatized block + // reductions + }; + + // Alias the temporary allocations from the single storage blob (or + // compute the necessary size of the blob) + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage + // allocation + return cudaSuccess; + } + + // Alias the allocation for the privatized per-block reductions + AccumT* d_block_reductions = static_cast(allocations[0]); + + // Get grid size for device_reduce_sweep_kernel + int reduce_grid_size = even_share.grid_size; + +// Log device_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceReduceKernel<<<%lu, %d, 0, %lld>>>(), %d items " + "per thread, %d SM occupancy\n", + (unsigned long) reduce_grid_size, + active_policy.Reduce().ThreadsPerBlock(), + (long long) stream, + active_policy.Reduce().ItemsPerThread(), + reduce_config.sm_occupancy); +#endif // CUB_DEBUG_LOG + + // Invoke DeviceReduceKernel + launcher_factory(reduce_grid_size, active_policy.Reduce().ThreadsPerBlock(), 0, stream) + .doit(reduce_kernel, d_in, d_block_reductions, num_items, even_share, reduction_op, init, transform_op); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + +// Log single_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceReduceSingleTileKernel<<<1, %d, 0, %lld>>>(), " + "%d items per thread\n", + active_policy.SingleTile().ThreadsPerBlock(), + (long long) stream, + active_policy.SingleTile().ItemsPerThread()); +#endif // CUB_DEBUG_LOG + + // Invoke DeviceReduceSingleTileKernel + launcher_factory(1, active_policy.SingleTile().ThreadsPerBlock(), 0, stream) + .doit( + single_tile_kernel, d_block_reductions, d_out, reduce_grid_size, reduction_op, init, ::cuda::std::identity{}); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + return CubDebug(detail::DebugSyncStream(stream)); + } + + //--------------------------------------------------------------------------- + // Chained policy invocation + //--------------------------------------------------------------------------- + + /// Invocation + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) + { + auto wrapped_policy = detail::reduce::MakeReducePolicyWrapper(active_policy); + if (num_items <= static_cast( + wrapped_policy.SingleTile().ThreadsPerBlock() * wrapped_policy.SingleTile().ItemsPerThread())) + { + // Small, single tile size + return InvokeSingleTile(kernel_source.SingleTileKernel(), wrapped_policy); + } + else + { + // Regular size + return InvokePasses(kernel_source.ReductionKernel(), kernel_source.SingleTileSecondKernel(), wrapped_policy); + } + } + + //--------------------------------------------------------------------------- + // Dispatch entrypoints + //--------------------------------------------------------------------------- + + /** + * @brief Internal dispatch routine for computing a device-wide reduction + * + * @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_in + * Pointer to the input sequence of data items + * + * @param[out] d_out + * Pointer to the output aggregate + * + * @param[in] num_items + * Total number of input items (i.e., length of `d_in`) + * + * @param[in] reduction_op + * Binary reduction functor + * + * @param[in] init + * The initial value of the reduction + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. + * Default is stream0. + */ + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + TransformOpT transform_op = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) + { + // Get PTX version + int ptx_version = 0; + if (const auto error = CubDebug(launcher_factory.PtxVersion(ptx_version))) + { + return error; + } + + // Create dispatch functor + DispatchReduce dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_items, + reduction_op, + init, + stream, + ptx_version, + transform_op, + kernel_source, + launcher_factory); + + // Ignore Wmaybe-uninitialized to work around a GCC 13 issue: + // https://github.com/NVIDIA/cccl/issues/4053 + _CCCL_DIAG_PUSH + _CCCL_DIAG_SUPPRESS_GCC("-Wmaybe-uninitialized") + // Dispatch to chained policy + return CubDebug(max_policy.Invoke(ptx_version, dispatch)); + _CCCL_DIAG_POP + } +}; + +/** + * @brief Utility class for dispatching the appropriately-tuned kernels for + * device-wide transform reduce + * + * Deprecated [Since 3.5] + * + * @tparam InputIteratorT + * Random-access input iterator type for reading input items @iterator + * + * @tparam OutputIteratorT + * Output iterator type for recording the reduced aggregate @iterator + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam ReductionOpT + * Binary reduction functor type having member + * `auto operator()(const T &a, const U &b)` + * + * @tparam TransformOpT + * Unary transform functor type having member + * `auto operator()(const T &a)` + * + * @tparam InitValueT + * Initial value type + */ +_CCCL_SUPPRESS_DEPRECATED_PUSH +template < + typename InputIteratorT, + typename OutputIteratorT, + typename OffsetT, + typename ReductionOpT, + typename TransformOpT, + typename InitValueT, + typename AccumT = + ::cuda::std::__accumulator_t>, + InitValueT>, + typename PolicyHub = detail::reduce::policy_hub, + typename KernelSource = detail::reduce::DeviceReduceKernelSource< + typename PolicyHub::MaxPolicy, + InputIteratorT, + OutputIteratorT, + OffsetT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT, + TransformOpT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +using DispatchTransformReduce CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce") = + DispatchReduce; +_CCCL_SUPPRESS_DEPRECATED_POP + +namespace detail::reduce +{ +// Retrieves a device pointer from a pointer-to-pointer. +// +// For CCCL.C's indirect_arg_t: ptr holds the address of the device pointer (&it.state). +// For regular C++ pointers: the caller passes &device_ptr directly. +// In both cases, dereferencing yields the actual device pointer. +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE void* get_device_ptr(void* ptr) +{ + return *reinterpret_cast(ptr); +} + +//! Preserve caller-selected immediate offset types; select a concrete offset type for deferred arguments. +template +using num_items_offset_t = + ::cuda::std::conditional_t<::cuda::args::__traits::is_deferred, + detail::choose_offset_t::element_type>, + typename ::cuda::args::__traits::element_type>; + +//! Creates the kernel argument for an immediate or deferred problem size without reading a deferred source. +//! Immediate values are cast to the selected offset type; deferred arguments are stripped to their source. +template +[[nodiscard]] CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE constexpr auto make_num_items_kernel_arg(OffsetT num_items) noexcept +{ + using args_traits_t = ::cuda::args::__traits; + using element_t = typename args_traits_t::element_type; + + if constexpr (args_traits_t::is_deferred) + { + static_assert(args_traits_t::is_single_value, "num_items must be a single value wrapped in cuda::args::deferred"); + static_assert(::cuda::std::__cccl_is_integer_v, "the num_items element type must be an integer"); + static_assert( + sizeof(element_t) == sizeof(::cuda::std::int32_t) || sizeof(element_t) == sizeof(::cuda::std::int64_t)); + } + + return CUB_NS_QUALIFIER::detail::parameter_from_host>(num_items); +} + +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_regular_size_reduce( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + TransformOpT transform_op, + ReducePolicy active_policy, + KernelSource kernel_source, + KernelLauncherFactory launcher_factory) +{ + using offset_t = num_items_offset_t; + + const auto kernel_num_items = make_num_items_kernel_arg(num_items); + + // Get SM count + int sm_count = 0; + if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) + { + return error; + } + + // Init regular kernel configuration + int sm_occupancy = 0; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + sm_occupancy, kernel_source.ReductionKernel(), active_policy.multi_tile.threads_per_block))) + { + return error; + } + + const int reduce_device_occupancy = sm_occupancy * sm_count; + const int max_blocks = reduce_device_occupancy * detail::subscription_factor; + + [[maybe_unused]] AccumT* d_block_reductions = nullptr; // buffer for per-block aggregates for the two-phase code path + if constexpr (!StableReductionOrder) + { + if (const auto error = + CubDebug(launcher_factory.MemsetAsync(get_device_ptr(&d_out), 0, kernel_source.InitSize(), stream))) + { + return error; + } + } + else + { + // Temporary storage allocation requirements + void* allocations[1] = {}; + size_t allocation_sizes[1] = { + max_blocks * kernel_source.AccumSize() // bytes needed for privatized block reductions + }; + + // Alias the temporary allocations from the single storage blob (or + // compute the necessary size of the blob) + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage allocation + return cudaSuccess; + } + + // Alias the allocation for the privatized per-block reductions + d_block_reductions = static_cast(allocations[0]); + } + + GridEvenShare even_share; + if constexpr (!::cuda::args::__traits::is_deferred) + { + const auto tile_size = active_policy.multi_tile.threads_per_block * active_policy.multi_tile.items_per_thread; + even_share.DispatchInit(kernel_num_items, max_blocks, tile_size); + } + + const int reduce_grid_size = [&] { + if constexpr (::cuda::args::__traits::is_deferred) + { + return max_blocks; + } + else if constexpr (!StableReductionOrder) + { + // The grid size for DeviceReduceKernel can be zero if the input size is zero. + // The atomic code path does not run a second kernel, so block zero handles an empty input. + return ::cuda::std::max(1, even_share.grid_size); + } + else + { + return even_share.grid_size; + } + }(); + +// Log device_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceReduceKernel<<<%lu, %d, 0, %lld>>>(), %d items " + "per thread, %d SM occupancy\n", + (unsigned long) reduce_grid_size, + active_policy.multi_tile.threads_per_block, + (long long) stream, + active_policy.multi_tile.items_per_thread, + sm_occupancy); +#endif // CUB_DEBUG_LOG + + // Invoke DeviceReduceKernel + auto reduce_kernel_output = [&] { + if constexpr (!StableReductionOrder) + { + return d_out; + } + else + { + return d_block_reductions; + } + }(); + if (const auto error = CubDebug( + launcher_factory(reduce_grid_size, active_policy.multi_tile.threads_per_block, 0, stream) + .doit(kernel_source.ReductionKernel(), + d_in, + reduce_kernel_output, + kernel_num_items, + even_share, + reduction_op, + init, + transform_op))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + if constexpr (StableReductionOrder) + { + // Log single_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceReduceSingleTileKernel<<<1, %d, 0, %lld>>>(), " + "%d items per thread\n", + active_policy.single_tile.threads_per_block, + (long long) stream, + active_policy.single_tile.items_per_thread); +#endif // CUB_DEBUG_LOG + + // Invoke DeviceReduceSingleTileKernel/DeviceReduceDeferredSingleTileKernel + if constexpr (::cuda::args::__traits::is_deferred) + { + if (const auto error = CubDebug( + launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(kernel_source.DeferredSingleTileSecondKernel(), + d_block_reductions, + d_out, + kernel_num_items, + reduce_grid_size, + reduction_op, + init, + ::cuda::std::identity{}))) + { + return error; + } + } + else + { + if (const auto error = CubDebug( + launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(kernel_source.SingleTileSecondKernel(), + d_block_reductions, + d_out, + reduce_grid_size, + reduction_op, + init, + ::cuda::std::identity{}))) + { + return error; + } + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} + +// select the accumulator type using an overload set, so __accumulator_t and invoke_result_t are not instantiated when +// an overriding accumulator type is present. This is needed by CCCL.C, which uses void as accumulator type. +template >> +_CCCL_HOST_DEVICE_API auto select_accum_t(use_default*) -> ::cuda::std::__accumulator_t< + ReductionOpT, + InputT, + ::cuda::std::conditional_t<::cuda::std::is_same_v, InputT, InitValueT>>; + +template , int> = 0> +_CCCL_HOST_DEVICE_API auto select_accum_t(OverrideAccumT*) -> OverrideAccumT; + +template < + typename OverrideAccumT = use_default, + bool StableReductionOrder = true, + typename InputIteratorT, + typename OutputIteratorT, + typename OffsetT, + typename ReductionOpT, + typename InitValueT = non_void_value_t>, + typename TransformOpT = ::cuda::std::identity, + typename AccumT = decltype(select_accum_t( + static_cast(nullptr))), + typename PolicySelector = policy_selector_from_types< + AccumT, + num_items_offset_t, + ReductionOpT, + StableReductionOrder ? __determinism_t::__run_to_run : __determinism_t::__not_guaranteed>, + typename KernelSource = DeviceReduceKernelSource< + PolicySelector, + InputIteratorT, + OutputIteratorT, + num_items_offset_t, + CUB_NS_QUALIFIER::detail::parameter_from_host_t, OffsetT>, + ReductionOpT, + InitValueT, + AccumT, + TransformOpT, + StableReductionOrder>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires reduce_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + TransformOpT transform_op = {}, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + using offset_t = num_items_offset_t; + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + // TODO: Remove this workaround once nvcc versions older than 12.4 are no longer supported. + // Older nvcc versions eagerly instantiate discarded statements in generic lambdas, so perform this conversion here. + // Both suppressions are needed for "never referenced" and "set but never used" diagnostics across supported nvcc + // and MSVC combinations. + [[maybe_unused]] offset_t offset_num_items{}; + if constexpr (StableReductionOrder && !::cuda::args::__traits::is_deferred) + { + offset_num_items = static_cast(num_items); + } + (void) offset_num_items; + + return dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) { + CUB_DETAIL_CONSTEXPR_ISH const ReducePolicy active_policy = policy_getter(); + + // known operators for integers are stable, even when using a non-deterministic reduction order + if constexpr (StableReductionOrder + && (!::cuda::std::is_integral_v || !is_cuda_binary_operator) ) + { + CUB_DETAIL_STATIC_ISH_ASSERT( + active_policy.multi_tile.reduce_algorithm != BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC, + "A run-to-run deterministic reduction must not use a non-deterministic reduce_algorithm"); + CUB_DETAIL_STATIC_ISH_ASSERT( + active_policy.single_tile.reduce_algorithm != BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC, + "A run-to-run deterministic reduction must not use a non-deterministic reduce_algorithm"); + } + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceReduce to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + if constexpr (StableReductionOrder && !::cuda::args::__traits::is_deferred) + { + const bool single_tile_problem = + offset_num_items <= (static_cast(active_policy.single_tile.threads_per_block) + * active_policy.single_tile.items_per_thread); + + // if the problem is small enough to fit into a single tile, just handle it and return early + if (single_tile_problem) + { + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceReduceSingleTileKernel<<<1, %d, 0, %lld>>>(), " + "%d items per thread\n", + active_policy.single_tile.threads_per_block, + (long long) stream, + active_policy.single_tile.items_per_thread); +#endif // CUB_DEBUG_LOG + + // Invoke single_reduce_sweep_kernel + if (const auto error = CubDebug( + launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit( + kernel_source.SingleTileKernel(), d_in, d_out, offset_num_items, reduction_op, init, transform_op))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + return cudaSuccess; + } + } + else if constexpr (!StableReductionOrder) + { + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + } + + // Regular size + return invoke_regular_size_reduce( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_items, + reduction_op, + init, + stream, + transform_op, + active_policy, + kernel_source, + launcher_factory); + }); +} +} // namespace detail::reduce + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce_by_key.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce_by_key.cuh new file mode 100644 index 00000000..91560416 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce_by_key.cuh @@ -0,0 +1,862 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2022, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * @brief cub::DeviceReduceByKey provides device-wide, parallel operations for + * reducing segments of values residing within device-accessible memory. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Kernel entry points + *****************************************************************************/ + +namespace detail::reduce_by_key +{ +template +struct streaming_context +{ + bool first_partition; + bool last_partition; + PrecedingKeyItT preceding_key_it; + + // We use a double-buffer to track the aggregate of the last run of the previous partition + AccumT* preceding_prefix; + AccumT* prefix_out; + + // We use a double-buffer to track the number of runs of previous partition + GlobalOffsetT* d_num_previous_uniques_in; + GlobalOffsetT* d_num_accumulated_uniques_out; + + _CCCL_DEVICE _CCCL_FORCEINLINE GlobalOffsetT num_accumulated_uniques_out() const + { + return first_partition ? GlobalOffsetT{0} : *d_num_previous_uniques_in; + }; + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE bool is_first_partition() const + { + return first_partition; + } + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE bool is_last_partition() const + { + return last_partition; + } + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE auto predecessor_key() const + { + return *preceding_key_it; + } + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE AccumT prefix() const + { + return *preceding_prefix; + } + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE void write_prefix(AccumT prefix) const + { + *prefix_out = prefix; + } + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE GlobalOffsetT* previous_uniques_ptr() const + { + return d_num_previous_uniques_in; + } + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE GlobalOffsetT num_uniques() const + { + return num_accumulated_uniques_out(); + } + + template + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE GlobalOffsetT add_num_uniques(NumUniquesT num_uniques) const + { + GlobalOffsetT total_uniques = num_accumulated_uniques_out() + static_cast(num_uniques); + + // The double buffer is only read by a subsequent partition and is not allocated for single-partition invocations + if (!last_partition) + { + *d_num_accumulated_uniques_out = total_uniques; + } + + return total_uniques; + } +}; + +/** + * @brief Multi-block reduce-by-key sweep kernel entry point + * + * @tparam PolicySelector + * Selects the tuning policy + * + * @tparam KeysInputIteratorT + * Random-access input iterator type for keys + * + * @tparam UniqueOutputIteratorT + * Random-access output iterator type for keys + * + * @tparam ValuesInputIteratorT + * Random-access input iterator type for values + * + * @tparam AggregatesOutputIteratorT + * Random-access output iterator type for values + * + * @tparam NumRunsOutputIteratorT + * Output iterator type for recording number of segments encountered + * + * @tparam ScanTileStateT + * Tile status interface type + * + * @tparam EqualityOpT + * KeyT equality operator type + * + * @tparam ReductionOpT + * ValueT reduction operator type + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @param d_keys_in + * Pointer to the input sequence of keys + * + * @param d_unique_out + * Pointer to the output sequence of unique keys (one key per run) + * + * @param d_values_in + * Pointer to the input sequence of corresponding values + * + * @param d_aggregates_out + * Pointer to the output sequence of value aggregates (one aggregate per run) + * + * @param d_num_runs_out + * Pointer to total number of runs encountered + * (i.e., the length of d_unique_out) + * + * @param tile_state + * Tile status interface + * + * @param start_tile + * The starting tile for the current grid + * + * @param equality_op + * KeyT equality operator + * + * @param reduction_op + * ValueT reduction operator + * + * @param num_items + * Total number of items to select from + */ +template +#if _CCCL_HAS_CONCEPTS() + requires reduce_by_key_policy_selector +#endif +__launch_bounds__(int(current_policy().lookback.threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceReduceByKeyKernel( + const KeysInputIteratorT d_keys_in, + const UniqueOutputIteratorT d_unique_out, + const ValuesInputIteratorT d_values_in, + const AggregatesOutputIteratorT d_aggregates_out, + const NumRunsOutputIteratorT d_num_runs_out, + ScanTileStateT tile_state, + const int start_tile, + EqualityOpT equality_op, + ReductionOpT reduction_op, + const OffsetT num_items, + const StreamingContextT streaming_context, + vsmem_t vsmem) +{ + static constexpr ReduceByKeyPolicy policy = current_policy(); + using AgentReduceByKeyPolicyT = agent_reduce_by_key_policy< + policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + delay_constructor_t>; + + using vsmem_helper_t = vsmem_helper_default_fallback_policy_t< + AgentReduceByKeyPolicyT, + AgentReduceByKey, + KeysInputIteratorT, + UniqueOutputIteratorT, + ValuesInputIteratorT, + AggregatesOutputIteratorT, + NumRunsOutputIteratorT, + EqualityOpT, + ReductionOpT, + OffsetT, + AccumT, + StreamingContextT>; + + // Thread block type for reducing tiles of value segments + using agent_reduce_by_key_t = typename vsmem_helper_t::agent_t; + + // Static shared memory allocation + __shared__ typename vsmem_helper_t::static_temp_storage_t static_temp_storage; + + // Get temporary storage + typename agent_reduce_by_key_t::TempStorage& temp_storage = + vsmem_helper_t::get_temp_storage(static_temp_storage, vsmem); + + // Process tiles + agent_reduce_by_key_t( + temp_storage, + d_keys_in, + d_unique_out, + d_values_in, + d_aggregates_out, + d_num_runs_out, + equality_op, + reduction_op, + streaming_context) + .ConsumeRange(num_items, tile_state, start_tile); + + // If applicable, hints to discard modified cache lines for vsmem + vsmem_helper_t::discard_temp_storage(temp_storage); +} +} // namespace detail::reduce_by_key + +/****************************************************************************** + * Dispatch + ******************************************************************************/ + +/** + * @brief Utility class for dispatching the appropriately-tuned kernels for + * DeviceReduceByKey + * + * Deprecated [Since 3.5] + * + * @tparam KeysInputIteratorT + * Random-access input iterator type for keys + * + * @tparam UniqueOutputIteratorT + * Random-access output iterator type for keys + * + * @tparam ValuesInputIteratorT + * Random-access input iterator type for values + * + * @tparam AggregatesOutputIteratorT + * Random-access output iterator type for values + * + * @tparam NumRunsOutputIteratorT + * Output iterator type for recording number of segments encountered + * + * @tparam EqualityOpT + * KeyT equality operator type + * + * @tparam ReductionOpT + * ValueT reduction operator type + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam PolicyHub + * Implementation detail, do not specify directly, requirements on the + * content of this type are subject to breaking change. + */ +template , + cub::detail::it_value_t>, + typename PolicyHub = detail::reduce_by_key::policy_hub< + ReductionOpT, + AccumT, + cub::detail::non_void_value_t>>> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceReduce::ReduceByKey") DispatchReduceByKey +{ + //------------------------------------------------------------------------- + // Types and constants + //------------------------------------------------------------------------- + + // The input values type + using ValueInputT = cub::detail::it_value_t; + + // Type used to provide context for streaming invocations (this is currently not used by ReduceByKey yet) + using streaming_context_t = NullType; + + static constexpr int INIT_KERNEL_THREADS = 128; + + // Tile status descriptor interface type + using ScanTileStateT = ReduceByKeyScanTileState; + + void* d_temp_storage; + size_t& temp_storage_bytes; + KeysInputIteratorT d_keys_in; + UniqueOutputIteratorT d_unique_out; + ValuesInputIteratorT d_values_in; + AggregatesOutputIteratorT d_aggregates_out; + NumRunsOutputIteratorT d_num_runs_out; + EqualityOpT equality_op; + ReductionOpT reduction_op; + OffsetT num_items; + cudaStream_t stream; + + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchReduceByKey( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + UniqueOutputIteratorT d_unique_out, + ValuesInputIteratorT d_values_in, + AggregatesOutputIteratorT d_aggregates_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + ReductionOpT reduction_op, + OffsetT num_items, + cudaStream_t stream) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_keys_in(d_keys_in) + , d_unique_out(d_unique_out) + , d_values_in(d_values_in) + , d_aggregates_out(d_aggregates_out) + , d_num_runs_out(d_num_runs_out) + , equality_op(equality_op) + , reduction_op(reduction_op) + , num_items(num_items) + , stream(stream) + {} + + //--------------------------------------------------------------------- + // Dispatch entrypoints + //--------------------------------------------------------------------- + + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t + Invoke(ScanInitKernelT init_kernel, ReduceByKeyKernelT reduce_by_key_kernel) + { + using vsmem_helper_t = detail::vsmem_helper_default_fallback_policy_t< + typename ActivePolicyT::ReduceByKeyPolicyT, + detail::reduce_by_key::AgentReduceByKey, + KeysInputIteratorT, + UniqueOutputIteratorT, + ValuesInputIteratorT, + AggregatesOutputIteratorT, + NumRunsOutputIteratorT, + EqualityOpT, + ReductionOpT, + OffsetT, + AccumT, + streaming_context_t>; + + constexpr int threads_per_block = vsmem_helper_t::agent_policy_t::BLOCK_THREADS; + constexpr int items_per_thread = vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD; + + cudaError error = cudaSuccess; + do + { + // Get device ordinal + int device_ordinal; + error = CubDebug(cudaGetDevice(&device_ordinal)); + if (cudaSuccess != error) + { + break; + } + + // Number of input tiles + int tile_size = threads_per_block * items_per_thread; + int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + // The amount of virtual shared memory to allocate + const auto vsmem_size = num_tiles * vsmem_helper_t::vsmem_per_block; + + // Specify temporary storage allocation requirements + size_t tile_descriptor_memory{}; + error = CubDebug(ScanTileStateT::AllocationSize(num_tiles, tile_descriptor_memory)); + if (cudaSuccess != error) + { + break; // bytes needed for tile status descriptors + } + size_t allocation_sizes[2] = {tile_descriptor_memory, vsmem_size}; + + // Compute allocation pointers into the single storage blob (or compute + // the necessary size of the blob) + void* allocations[2] = {}; + + error = CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)); + if (cudaSuccess != error) + { + break; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage + // allocation + break; + } + + // Construct the tile status interface + ScanTileStateT tile_state; + error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0])); + if (cudaSuccess != error) + { + break; + } + + // Log init_kernel configuration + int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(num_tiles, INIT_KERNEL_THREADS)); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, INIT_KERNEL_THREADS, (long long) stream); +#endif // CUB_DEBUG_LOG + + // Invoke init_kernel to initialize tile descriptors + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, INIT_KERNEL_THREADS, 0, stream) + .doit(init_kernel, tile_state, num_tiles, d_num_runs_out)); + if (cudaSuccess != error) + { + break; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(detail::DebugSyncStream(stream)); + if (cudaSuccess != error) + { + break; + } + + // Return if empty problem: note, we're initializing d_num_runs_out to 0 in init_kernel above + if (num_items == 0) + { + break; + } + + // Get SM occupancy for reduce_by_key_kernel + int reduce_by_key_sm_occupancy; + error = CubDebug(MaxSmOccupancy(reduce_by_key_sm_occupancy, reduce_by_key_kernel, threads_per_block)); + + if (cudaSuccess != error) + { + break; + } + + // Get max x-dimension of grid + int max_dim_x; + error = CubDebug(cudaDeviceGetAttribute(&max_dim_x, cudaDevAttrMaxGridDimX, device_ordinal)); + if (cudaSuccess != error) + { + break; + } + + // Run grids in epochs (in case number of tiles exceeds max x-dimension + int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x); + for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size) + { +// Log reduce_by_key_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking %d reduce_by_key_kernel<<<%d, %d, 0, %lld>>>(), %d " + "items per thread, %d SM occupancy\n", + start_tile, + scan_grid_size, + threads_per_block, + (long long) stream, + items_per_thread, + reduce_by_key_sm_occupancy); +#endif // CUB_DEBUG_LOG + + // Invoke reduce_by_key_kernel + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(scan_grid_size, threads_per_block, 0, stream) + .doit(reduce_by_key_kernel, + d_keys_in, + d_unique_out, + d_values_in, + d_aggregates_out, + d_num_runs_out, + tile_state, + start_tile, + equality_op, + reduction_op, + num_items, + streaming_context_t{}, + cub::detail::vsmem_t{allocations[1]})); + if (cudaSuccess != error) + { + break; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(detail::DebugSyncStream(stream)); + if (cudaSuccess != error) + { + break; + } + } + } while (false); + + return error; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke() + { + return Invoke( + detail::scan::DeviceCompactInitKernel, + detail::reduce_by_key::DeviceReduceByKeyKernel< + detail::reduce_by_key::policy_selector_from_hub, + KeysInputIteratorT, + UniqueOutputIteratorT, + ValuesInputIteratorT, + AggregatesOutputIteratorT, + NumRunsOutputIteratorT, + ScanTileStateT, + EqualityOpT, + ReductionOpT, + OffsetT, + AccumT, + streaming_context_t>); + } + + /** + * Internal dispatch routine + * @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 + * Pointer to the input sequence of keys + * + * @param[out] d_unique_out + * Pointer to the output sequence of unique keys (one key per run) + * + * @param[in] d_values_in + * Pointer to the input sequence of corresponding values + * + * @param[out] d_aggregates_out + * Pointer to the output sequence of value aggregates + * (one aggregate per run) + * + * @param[out] d_num_runs_out + * Pointer to total number of runs encountered + * (i.e., the length of d_unique_out) + * + * @param[in] equality_op + * KeyT equality operator + * + * @param[in] reduction_op + * ValueT reduction operator + * + * @param[in] num_items + * Total number of items to select from + * + * @param[in] stream + * CUDA stream to launch kernels within. Default is stream0. + */ + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + UniqueOutputIteratorT d_unique_out, + ValuesInputIteratorT d_values_in, + AggregatesOutputIteratorT d_aggregates_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + ReductionOpT reduction_op, + OffsetT num_items, + cudaStream_t stream) + { + cudaError error = cudaSuccess; + + do + { + // Get PTX version + int ptx_version = 0; + error = CubDebug(PtxVersion(ptx_version)); + if (cudaSuccess != error) + { + break; + } + + DispatchReduceByKey dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_unique_out, + d_values_in, + d_aggregates_out, + d_num_runs_out, + equality_op, + reduction_op, + num_items, + stream); + + // Dispatch + error = CubDebug(PolicyHub::MaxPolicy::Invoke(ptx_version, dispatch)); + if (cudaSuccess != error) + { + break; + } + } while (false); + + return error; + } +}; + +namespace detail::reduce_by_key +{ +// we move the conversion of the policy to the agent policy and its use out of the lambda below, so MSVC does not ICE +template +_CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_getter) +{ + // TODO(bgruber): refactor this in the future + constexpr ReduceByKeyPolicy policy = policy_getter(); + using Policy = agent_reduce_by_key_policy< + policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + delay_constructor_t>; + using vsmem_helper_t = vsmem_helper_default_fallback_policy_t; + return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS, + vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD, + vsmem_helper_t::vsmem_per_block}; +} + +template >, + typename KeyT = non_void_value_t>, + typename PolicySelector = policy_selector_from_types> +#if _CCCL_HAS_CONCEPTS() + requires reduce_by_key::reduce_by_key_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + UniqueOutputIteratorT d_unique_out, + ValuesInputIteratorT d_values_in, + AggregatesOutputIteratorT d_aggregates_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + ReductionOpT reduction_op, + OffsetT num_items, + cudaStream_t stream, + PolicySelector policy_selector = {}) +{ + using streaming_context_t = NullType; // streaming context not used for ReduceByKey yet + using ScanTileStateT = ReduceByKeyScanTileState; + [[maybe_unused]] static constexpr int init_kernel_threads = 128; + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(ptx_compute_cap(cc))) + { + return error; + } + + return detail::dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) { +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << policy_getter(); + _CubLog("Dispatching DeviceReduceByKey to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const auto [threads_per_block, items_per_thread, vsmem_per_block] = determine_threads_items_vsmem< + decltype(policy_getter), + KeysInputIteratorT, + UniqueOutputIteratorT, + ValuesInputIteratorT, + AggregatesOutputIteratorT, + NumRunsOutputIteratorT, + EqualityOpT, + ReductionOpT, + OffsetT, + AccumT, + streaming_context_t>(policy_getter); + + // Number of input tiles + const int tile_size = threads_per_block * items_per_thread; + const int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + // The amount of virtual shared memory to allocate + const auto vsmem_size = num_tiles * vsmem_per_block; + + size_t tile_descriptor_memory{}; + if (const auto error = CubDebug(ScanTileStateT::AllocationSize(num_tiles, tile_descriptor_memory))) + { + return error; + } + size_t allocation_sizes[2] = {tile_descriptor_memory, vsmem_size}; + void* allocations[2] = {}; + + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + ScanTileStateT tile_state; + if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + const int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(num_tiles, init_kernel_threads)); +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, init_kernel_threads, (long long) stream); +#endif + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_kernel_threads, 0, stream) + .doit(detail::scan::DeviceCompactInitKernel, + tile_state, + num_tiles, + d_num_runs_out))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + if (num_items == 0) + { + return cudaSuccess; + } + + auto reduce_by_key_kernel = &DeviceReduceByKeyKernel< + PolicySelector, + KeysInputIteratorT, + UniqueOutputIteratorT, + ValuesInputIteratorT, + AggregatesOutputIteratorT, + NumRunsOutputIteratorT, + ScanTileStateT, + EqualityOpT, + ReductionOpT, + OffsetT, + AccumT, + streaming_context_t>; + + int reduce_by_key_sm_occupancy{}; + if (const auto error = + CubDebug(MaxSmOccupancy(reduce_by_key_sm_occupancy, reduce_by_key_kernel, threads_per_block))) + { + return error; + } + + int device_ordinal{}; + if (const auto error = CubDebug(cudaGetDevice(&device_ordinal))) + { + return error; + } + int max_dim_x{}; + if (const auto error = CubDebug(cudaDeviceGetAttribute(&max_dim_x, cudaDevAttrMaxGridDimX, device_ordinal))) + { + return error; + } + + const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x); + for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size) + { +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking %d reduce_by_key_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread, %d SM occupancy\n", + start_tile, + scan_grid_size, + threads_per_block, + (long long) stream, + items_per_thread, + reduce_by_key_sm_occupancy); +#endif + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(scan_grid_size, threads_per_block, 0, stream) + .doit(reduce_by_key_kernel, + d_keys_in, + d_unique_out, + d_values_in, + d_aggregates_out, + d_num_runs_out, + tile_state, + start_tile, + equality_op, + reduction_op, + num_items, + streaming_context_t{}, + cub::detail::vsmem_t{allocations[1]}))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + return cudaSuccess; + }); +} +} // namespace detail::reduce_by_key + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce_deterministic.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce_deterministic.cuh new file mode 100644 index 00000000..49c8608b --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_reduce_deterministic.cuh @@ -0,0 +1,484 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +//! @file +//! This file device-wide, parallel operations for computing a reduction across a sequence of data items residing within +//! device-accessible memory. Current reduction operator supported is ``cuda::std::plus`` + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::rfa +{ +using cuda::execution::determinism::__determinism_t; +using reduce::policy_selector_from_types; + +template +using transformed_input_t = ::cuda::std::decay_t<::cuda::std::invoke_result_t>; + +template +using accum_t = ::cuda::std:: + __accumulator_t<::cuda::std::plus<>, InitValueT, transformed_input_t>>; + +template >* = nullptr> +struct deterministic_sum_t +{ + using DeterministicAcc = ReproducibleFloatingAccumulator; + + _CCCL_DEVICE DeterministicAcc operator()(DeterministicAcc acc, FloatType f) + { + acc += f; + return acc; + } + + _CCCL_DEVICE DeterministicAcc operator()(FloatType f, DeterministicAcc acc) + { + return this->operator()(acc, f); + } + + _CCCL_DEVICE DeterministicAcc operator()(DeterministicAcc lhs, DeterministicAcc rhs) + { + DeterministicAcc rtn = lhs; + rtn += rhs; + return rtn; + } + + _CCCL_DEVICE FloatType operator()(FloatType lhs, FloatType rhs) + { + return lhs + rhs; + } +}; + +template +CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invoke_single_tile( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + TransformOpT transform_op, + ReducePolicy active_policy, + KernelLauncherFactory launcher_factory) +{ + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + +// Log single_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeterministicDeviceReduceSingleTileKernel<<<1, %d, 0, %lld>>>(), " + "%d items per thread\n", + active_policy.single_tile.threads_per_block, + (long long) stream, + active_policy.single_tile.items_per_thread); +#endif // CUB_DEBUG_LOG + + // Invoke single_reduce_sweep_kernel + if (const auto error = CubDebug( + launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(detail::reduce::DeterministicDeviceReduceSingleTileKernel< + PolicySelector, + InputIteratorT, + OutputIteratorT, + ReductionOpT, + InitValueT, + DeterministicAccumT, + TransformOpT>, + d_in, + d_out, + static_cast(num_items), + reduction_op, + init, + transform_op))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + return CubDebug(detail::DebugSyncStream(stream)); +} + +template +CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t invoke_passes( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + TransformOpT transform_op, + ReducePolicy active_policy, + KernelLauncherFactory launcher_factory) +{ + // Immediate chunk sizes are passed to the kernels as-is; deferred problem sizes are read on device. + using num_items_kernel_t = detail::parameter_from_host_t; + + int sm_count; + if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) + { + return error; + } + + KernelConfig reduce_config; + if (const auto error = CubDebug(reduce_config.__init( + detail::reduce::DeterministicDeviceReduceKernel< + PolicySelector, + InputIteratorT, + num_items_kernel_t, + ReductionOpT, + DeterministicAccumT, + TransformOpT>, + active_policy.multi_tile))) + { + return error; + } + + const int reduce_device_occupancy = reduce_config.sm_occupancy * sm_count; + const int max_blocks = reduce_device_occupancy * detail::subscription_factor; + + const int num_items_per_chunk = ::cuda::std::numeric_limits<::cuda::std::int32_t>::max(); + + // A deferred problem size cannot be read on the host, so these defaults stand: the kernel consumes the whole + // problem in a single launch with the worst-case grid, whose surplus blocks exit early. + int num_chunks = 1; + int chunk_grid_size = max_blocks; + int partial_chunk_size = 0; + bool has_partial_chunk = false; + int last_chunk_grid_size = max_blocks; + if constexpr (!::cuda::args::__traits::is_deferred) + { + num_chunks = static_cast(::cuda::ceil_div(num_items, num_items_per_chunk)); + + const int chunk_tile_grid_size = ::cuda::ceil_div(num_items_per_chunk, reduce_config.tile_size); + chunk_grid_size = ::cuda::std::min(max_blocks, chunk_tile_grid_size); + + partial_chunk_size = num_items % num_items_per_chunk; + has_partial_chunk = partial_chunk_size != 0; + const int last_chunk_tile_grid_size = ::cuda::ceil_div(partial_chunk_size, reduce_config.tile_size); + + last_chunk_grid_size = ::cuda::std::min(max_blocks, last_chunk_tile_grid_size); + } + + const int reduce_grid_size = chunk_grid_size * (num_chunks - 1) + last_chunk_grid_size; + + // Temporary storage allocation requirements + void* allocations[1] = {}; + size_t allocation_sizes[1] = { + reduce_grid_size * sizeof(DeterministicAccumT) // bytes needed for privatized block reductions + }; + + // Alias the temporary allocations from the single storage blob (or + // compute the necessary size of the blob) + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage allocation + return cudaSuccess; + } + + // Alias the allocation for the privatized per-block reductions + DeterministicAccumT* d_block_reductions = static_cast(allocations[0]); + + auto d_chunk_block_reductions = d_block_reductions; + for (int chunk_index = 0; chunk_index < num_chunks; chunk_index++) + { + const int num_current_items = + ((chunk_index + 1 == num_chunks) && has_partial_chunk) ? partial_chunk_size : num_items_per_chunk; + + const auto current_grid_size = + static_cast(num_current_items == num_items_per_chunk ? chunk_grid_size : last_chunk_grid_size); + + // An immediate problem size is passed as the current chunk size; a deferred problem size is read on device and + // `num_current_items` holds the worst-case chunk size, which must not be passed to the kernel. + const auto kernel_num_items = [=] { + if constexpr (::cuda::args::__traits::is_deferred) + { + return detail::reduce::make_num_items_kernel_arg(num_items); + } + else + { + return num_current_items; + } + }(); + +// Log device_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeterministicDeviceReduceKernel<<<%d, %d, 0, %lld>>>(), %d items " + "per thread, %d SM occupancy\n", + current_grid_size, + active_policy.multi_tile.threads_per_block, + (long long) stream, + active_policy.multi_tile.items_per_thread, + reduce_config.sm_occupancy); +#endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + launcher_factory(current_grid_size, active_policy.multi_tile.threads_per_block, 0, stream) + .doit(detail::reduce::DeterministicDeviceReduceKernel, + d_in, + d_chunk_block_reductions, + kernel_num_items, + reduction_op, + transform_op, + current_grid_size))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + if (chunk_index + 1 < num_chunks) + { + d_in += num_current_items; + d_chunk_block_reductions += current_grid_size; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + +// Log single_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeterministicDeviceReduceSingleTileKernel<<<1, %d, 0, %lld>>>(), " + "%d items per thread\n", + active_policy.single_tile.threads_per_block, + (long long) stream, + active_policy.single_tile.items_per_thread); +#endif // CUB_DEBUG_LOG + + // Invoke DeterministicDeviceReduceSingleTileKernel/DeterministicDeviceReduceDeferredSingleTileKernel + const auto second_pass_error = [&] { + if constexpr (::cuda::args::__traits::is_deferred) + { + return launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(detail::reduce::DeterministicDeviceReduceDeferredSingleTileKernel< + PolicySelector, + DeterministicAccumT*, + OutputIteratorT, + num_items_kernel_t, + ReductionOpT, + InitValueT, + DeterministicAccumT>, + d_block_reductions, + d_out, + detail::reduce::make_num_items_kernel_arg(num_items), + reduce_grid_size, + reduction_op, + init, + ::cuda::std::identity{}); + } + else + { + return launcher_factory(1, active_policy.single_tile.threads_per_block, 0, stream) + .doit(detail::reduce::DeterministicDeviceReduceSingleTileKernel< + PolicySelector, + DeterministicAccumT*, + OutputIteratorT, + ReductionOpT, + InitValueT, + DeterministicAccumT>, + d_block_reductions, + d_out, + reduce_grid_size, + reduction_op, + init, + ::cuda::std::identity{}); + } + }(); + if (const auto error = CubDebug(second_pass_error)) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + return CubDebug(detail::DebugSyncStream(stream)); +} + +template , + typename PolicySelector = policy_selector_from_types, + deterministic_sum_t, + __determinism_t::__gpu_to_gpu>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + InitValueT init = {}, + cudaStream_t stream = {}, + TransformOpT transform_op = {}, + PolicySelector policy_selector = {}, + KernelLauncherFactory launcher_factory = {}) +{ + // Get CC + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + const ReducePolicy active_policy = policy_selector(cc); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceReduceDeterministic to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + using deterministic_add_t = deterministic_sum_t; + using input_unwrapped_it_t = THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t; + + input_unwrapped_it_t d_in_unwrapped = THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in); + + // A deferred problem size cannot be compared against the single-tile capacity on the host, so a deferred + // reduction always takes the two-pass path. + if constexpr (!::cuda::args::__traits::is_deferred) + { + const auto tile_items = static_cast(active_policy.single_tile.threads_per_block) + * static_cast(active_policy.single_tile.items_per_thread); + + if (num_items <= tile_items) + { + return invoke_single_tile( + d_temp_storage, + temp_storage_bytes, + d_in_unwrapped, + d_out, + num_items, + deterministic_add_t{}, + init, + stream, + transform_op, + active_policy, + launcher_factory); + } + } + + return invoke_passes( + d_temp_storage, + temp_storage_bytes, + d_in_unwrapped, + d_out, + num_items, + deterministic_add_t{}, + init, + stream, + transform_op, + active_policy, + launcher_factory); +} +} // namespace detail::rfa +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_rle.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_rle.cuh new file mode 100644 index 00000000..7592a836 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_rle.cuh @@ -0,0 +1,840 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::DeviceRle provides device-wide, parallel operations for run-length-encoding sequences of + * data items residing within device-accessible memory. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Kernel entry points + *****************************************************************************/ + +namespace detail::rle +{ +template +struct streaming_context +{ + bool first_partition; + bool last_partition; + // Global offset of the current partition to compute and write out the correct absolute offsets to the user + GlobalOffsetT current_base_offset; + + // We use a double-buffer to track the aggregated run-length of the last run of the previous partition + RunLengthT* preceding_length; + RunLengthT* length_out; + + // We use a double-buffer to track the number of runs of previous partition + GlobalOffsetT* d_num_previous_uniques_in; + GlobalOffsetT* d_num_accumulated_uniques_out; + + _CCCL_DEVICE _CCCL_FORCEINLINE GlobalOffsetT num_accumulated_uniques_out() const + { + return first_partition ? GlobalOffsetT{0} : *d_num_previous_uniques_in; + }; + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE RunLengthT prefix() const + { + return *preceding_length; + } + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE void write_prefix(RunLengthT prefix) const + { + *length_out = prefix; + } + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE GlobalOffsetT* previous_uniques_ptr() const + { + return d_num_previous_uniques_in; + } + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE GlobalOffsetT num_uniques() const + { + return num_accumulated_uniques_out(); + } + + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE GlobalOffsetT base_offset() const + { + return current_base_offset; + } + + template + _CCCL_FORCEINLINE _CCCL_HOST_DEVICE GlobalOffsetT add_num_uniques(NumUniquesT num_uniques) const + { + GlobalOffsetT total_uniques = num_accumulated_uniques_out() + static_cast(num_uniques); + + // The double buffer is only read by a subsequent partition and is not allocated for single-partition invocations + if (!last_partition) + { + *d_num_accumulated_uniques_out = total_uniques; + } + + return total_uniques; + } +}; + +/** + * Select kernel entry point (multi-block) + * + * Performs functor-based selection if SelectOp functor type != NullType + * Otherwise performs flag-based selection if FlagIterator's value type != NullType + * Otherwise performs discontinuity selection (keep unique) + * + * @tparam PolicySelector + * Selects the tuning policy + * + * @tparam InputIteratorT + * Random-access input iterator type for reading input items @iterator + * + * @tparam OffsetsOutputIteratorT + * Random-access output iterator type for writing run-offset values @iterator + * + * @tparam LengthsOutputIteratorT + * Random-access output iterator type for writing run-length values @iterator + * + * @tparam NumRunsOutputIteratorT + * Output iterator type for recording the number of runs encountered @iterator + * + * @tparam ScanTileStateT + * Tile status interface type + * + * @tparam EqualityOpT + * T equality operator type + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @param d_in + * Pointer to input sequence of data items + * + * @param d_offsets_out + * Pointer to output sequence of run-offsets + * + * @param d_lengths_out + * Pointer to output sequence of run-lengths + * + * @param d_num_runs_out + * Pointer to total number of runs (i.e., length of `d_offsets_out`) + * + * @param tile_status + * Tile status interface + * + * @param equality_op + * Equality operator for input items + * + * @param num_items + * Total number of input items (i.e., length of `d_in`) + * + * @param num_tiles + * Total number of tiles for the entire problem + */ +template +#if _CCCL_HAS_CONCEPTS() + requires non_trivial_runs::rle_non_trivial_runs_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().lookback.threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceRleSweepKernel( + const InputIteratorT d_in, + const OffsetsOutputIteratorT d_offsets_out, + const LengthsOutputIteratorT d_lengths_out, + const NumRunsOutputIteratorT d_num_runs_out, + ScanTileStateT tile_status, + EqualityOpT equality_op, + const OffsetT num_items, + const int num_tiles, + const StreamingContextT streaming_context) +{ + static constexpr RleNonTrivialRunsPolicy policy = current_policy(); + using AgentRlePolicyT = + agent_rle_policy>; + + using AgentRleT = + AgentRle; + + // Shared memory for AgentRle + __shared__ typename AgentRleT::TempStorage temp_storage; + + // Process tiles + AgentRleT(temp_storage, d_in, d_offsets_out, d_lengths_out, equality_op, num_items, streaming_context) + .ConsumeRange(num_tiles, tile_status, d_num_runs_out); +} + +// TODO(bgruber): remove in CCCL 4.0 when we drop the RLE dispatchers +template +struct policy_selector_from_hub +{ + [[nodiscard]] _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability /*cc*/) const + -> RleNonTrivialRunsPolicy + { + using RleSweepPolicyT = typename PolicyHub::MaxPolicy::RleSweepPolicyT; + return RleNonTrivialRunsPolicy{ + RleNonTrivialRunsAlgorithm::lookback, + { + RleSweepPolicyT::BLOCK_THREADS, + RleSweepPolicyT::ITEMS_PER_THREAD, + RleSweepPolicyT::LOAD_ALGORITHM, + RleSweepPolicyT::LOAD_MODIFIER, + RleSweepPolicyT::STORE_WARP_TIME_SLICING, + RleSweepPolicyT::SCAN_ALGORITHM, + lookback_delay_policy_from_type, + }, + }; + } +}; +} // namespace detail::rle + +/****************************************************************************** + * Dispatch + ******************************************************************************/ + +/** + * Utility class for dispatching the appropriately-tuned kernels for DeviceRle + * + * @tparam InputIteratorT + * Random-access input iterator type for reading input items @iterator + * + * @tparam OffsetsOutputIteratorT + * Random-access output iterator type for writing run-offset values @iterator + * + * @tparam LengthsOutputIteratorT + * Random-access output iterator type for writing run-length values @iterator + * + * @tparam NumRunsOutputIteratorT + * Output iterator type for recording the number of runs encountered @iterator + * + * @tparam EqualityOpT + * T equality operator type + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam PolicyHub + * Implementation detail, do not specify directly, requirements on the + * content of this type are subject to breaking change. + */ +template , + cub::detail::it_value_t>> +struct CCCL_DEPRECATED_BECAUSE("Please use DeviceRunLengthEncode") DeviceRleDispatch +{ + /****************************************************************************** + * Types and constants + ******************************************************************************/ + // Offsets to index items within one partition (i.e., a single kernel invocation) + using local_offset_t = ::cuda::std::int32_t; + + // If the number of items provided by the user may exceed the maximum number of items processed by a single kernel + // invocation, we may require multiple kernel invocations + static constexpr bool use_streaming_invocation = ::cuda::std::numeric_limits::max() + > ::cuda::std::numeric_limits::max(); + + // Offsets to index any item within the entire input (large enough to cover num_items) + using global_offset_t = OffsetT; + + // The lengths output value type + using length_t = cub::detail::non_void_value_t; + + // Type used to provide context about the current partition during a streaming invocation + using streaming_context_t = + ::cuda::std::conditional_t, + NullType>; + + static constexpr int init_kernel_threads = 128; + + // Tile status descriptor interface type + using ScanTileStateT = ReduceByKeyScanTileState; + + void* d_temp_storage; + size_t& temp_storage_bytes; + InputIteratorT d_in; + OffsetsOutputIteratorT d_offsets_out; + LengthsOutputIteratorT d_lengths_out; + NumRunsOutputIteratorT d_num_runs_out; + EqualityOpT equality_op; + global_offset_t num_items; + cudaStream_t stream; + + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DeviceRleDispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OffsetsOutputIteratorT d_offsets_out, + LengthsOutputIteratorT d_lengths_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + global_offset_t num_items, + cudaStream_t stream) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_in(d_in) + , d_offsets_out(d_offsets_out) + , d_lengths_out(d_lengths_out) + , d_num_runs_out(d_num_runs_out) + , equality_op(equality_op) + , num_items(num_items) + , stream(stream) + {} + + /****************************************************************************** + * Dispatch entrypoints + ******************************************************************************/ + + /** + * Internal dispatch routine for computing a device-wide run-length-encode using the + * specified kernel functions. + * + * @tparam DeviceScanInitKernelPtr + * Function type of cub::DeviceScanInitKernel + * + * @tparam DeviceRleSweepKernelPtr + * Function type of cub::DeviceRleSweepKernelPtr + * + * @param device_scan_init_kernel + * Kernel function pointer to parameterization of cub::DeviceScanInitKernel + * + * @param device_rle_sweep_kernel + * Kernel function pointer to parameterization of cub::DeviceRleSweepKernel + */ + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t + Invoke(DeviceScanInitKernelPtr device_scan_init_kernel, DeviceRleSweepKernelPtr device_rle_sweep_kernel) + { + cudaError error = cudaSuccess; + + constexpr int threads_per_block = ActivePolicyT::RleSweepPolicyT::BLOCK_THREADS; + constexpr int items_per_thread = ActivePolicyT::RleSweepPolicyT::ITEMS_PER_THREAD; + constexpr auto tile_size = static_cast(threads_per_block * items_per_thread); + + // The upper bound of for the number of items that a single kernel invocation will ever process + auto capped_num_items_per_invocation = num_items; + if constexpr (use_streaming_invocation) + { + capped_num_items_per_invocation = + static_cast(::cuda::std::numeric_limits::max()); + // Make sure that the number of items is a multiple of tile size + capped_num_items_per_invocation -= (capped_num_items_per_invocation % tile_size); + } + + // Across invocations, the maximum number of items that a single kernel invocation will ever process + const auto max_num_items_per_invocation = + use_streaming_invocation ? ::cuda::std::min(capped_num_items_per_invocation, num_items) : num_items; + + // Number of invocations required to "iterate" over the total input (at least one iteration to process zero items) + auto const num_partitions = + (capped_num_items_per_invocation == 0) + ? global_offset_t{1} + : ::cuda::ceil_div(num_items, capped_num_items_per_invocation); + + // Number of input tiles + int max_num_tiles = static_cast(::cuda::ceil_div(max_num_items_per_invocation, tile_size)); + + // Specify temporary storage allocation requirements + size_t allocation_sizes[3]; + error = CubDebug(ScanTileStateT::AllocationSize(max_num_tiles, allocation_sizes[0])); + if (cudaSuccess != error) + { + return error; + } + allocation_sizes[1] = num_partitions > 1 ? sizeof(global_offset_t) * 2 : size_t{0}; + allocation_sizes[2] = num_partitions > 1 ? sizeof(local_offset_t) * 2 : size_t{0}; + + // Compute allocation pointers into the single storage blob (or compute the necessary size of + // the blob) + void* allocations[3] = {}; + + error = CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)); + if (error != cudaSuccess) + { + return error; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage allocation + return error; + } + + // Iterate over the partitions until all input is processed + for (global_offset_t partition_idx = 0; partition_idx < num_partitions; partition_idx++) + { + global_offset_t current_partition_offset = partition_idx * capped_num_items_per_invocation; + global_offset_t current_num_items = + (partition_idx + 1 == num_partitions) + ? (num_items - current_partition_offset) + : capped_num_items_per_invocation; + + // Construct the tile status interface + const auto num_current_tiles = static_cast(::cuda::ceil_div(current_num_items, tile_size)); + + // Construct the tile status interface + ScanTileStateT tile_status; + error = CubDebug(tile_status.Init(num_current_tiles, allocations[0], allocation_sizes[0])); + if (cudaSuccess != error) + { + return error; + } + + // Log init_kernel configuration + int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(num_current_tiles, init_kernel_threads)); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking device_scan_init_kernel<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + init_kernel_threads, + (long long) stream); +#endif // CUB_DEBUG_LOG + + // Invoke device_scan_init_kernel to initialize tile descriptors and queue descriptors + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_kernel_threads, 0, stream) + .doit(device_scan_init_kernel, tile_status, num_current_tiles, d_num_runs_out)); + if (cudaSuccess != error) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(detail::DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + + // Return if empty problem: note, we're initializing d_num_runs_out to 0 in device_scan_init_kernel above + if (num_items <= 1) + { + return error; + } + +// Log device_rle_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking device_rle_sweep_kernel<<<%d, %d, 0, %lld>>>(), %d items per " + "thread\n", + num_current_tiles, + threads_per_block, + (long long) stream, + items_per_thread); +#endif // CUB_DEBUG_LOG + + // Invoke device_rle_sweep_kernel + if constexpr (use_streaming_invocation) + { + auto tmp_num_uniques = static_cast(allocations[1]); + auto tmp_prefix = static_cast(allocations[2]); + + const bool is_first_partition = (partition_idx == 0); + const bool is_last_partition = (partition_idx + 1 == num_partitions); + const int buffer_selector = partition_idx % 2; + + streaming_context_t streaming_context{ + is_first_partition, + is_last_partition, + current_partition_offset, + &tmp_prefix[buffer_selector], + &tmp_prefix[buffer_selector ^ 0x01], + &tmp_num_uniques[buffer_selector], + &tmp_num_uniques[buffer_selector ^ 0x01]}; + + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_current_tiles, threads_per_block, 0, stream) + .doit(device_rle_sweep_kernel, + d_in + current_partition_offset, + d_offsets_out, + d_lengths_out, + d_num_runs_out, + tile_status, + equality_op, + static_cast(current_num_items), + num_current_tiles, + streaming_context)); + if (cudaSuccess != error) + { + return error; + } + } + else + { + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_current_tiles, threads_per_block, 0, stream) + .doit(device_rle_sweep_kernel, + d_in + current_partition_offset, + d_offsets_out, + d_lengths_out, + d_num_runs_out, + tile_status, + equality_op, + static_cast(current_num_items), + num_current_tiles, + NullType{})); + if (cudaSuccess != error) + { + return error; + } + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(detail::DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + } + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke() + { + return Invoke( + detail::scan::DeviceCompactInitKernel, + detail::rle::DeviceRleSweepKernel< + detail::rle::policy_selector_from_hub, + InputIteratorT, + OffsetsOutputIteratorT, + LengthsOutputIteratorT, + NumRunsOutputIteratorT, + ScanTileStateT, + EqualityOpT, + local_offset_t, + global_offset_t, + streaming_context_t>); + } + + /** + * Internal dispatch routine + * + * @param 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 temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param d_in + * Pointer to input sequence of data items + * + * @param d_offsets_out + * Pointer to output sequence of run-offsets + * + * @param d_lengths_out + * Pointer to output sequence of run-lengths + * + * @param d_num_runs_out + * Pointer to total number of runs (i.e., length of `d_offsets_out`) + * + * @param equality_op + * Equality operator for input items + * + * @param num_items + * Total number of input items (i.e., length of `d_in`) + * + * @param stream + * **[optional]** CUDA stream to launch kernels within. + * Default is stream0. + */ + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OffsetsOutputIteratorT d_offsets_out, + LengthsOutputIteratorT d_lengths_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + OffsetT num_items, + cudaStream_t stream) + { + cudaError error = cudaSuccess; + + // Get PTX version + int ptx_version = 0; + error = CubDebug(PtxVersion(ptx_version)); + if (cudaSuccess != error) + { + return error; + } + + DeviceRleDispatch dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_offsets_out, + d_lengths_out, + d_num_runs_out, + equality_op, + num_items, + stream); + + // Dispatch + error = CubDebug(PolicyHub::MaxPolicy::Invoke(ptx_version, dispatch)); + if (cudaSuccess != error) + { + return error; + } + return cudaSuccess; + } +}; + +namespace detail::rle +{ +template , + typename key_t = it_value_t, + typename PolicySelector = non_trivial_runs::policy_selector_from_types> +#if _CCCL_HAS_CONCEPTS() + requires non_trivial_runs::rle_non_trivial_runs_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OffsetsOutputIteratorT d_offsets_out, + LengthsOutputIteratorT d_lengths_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + OffsetT num_items, + cudaStream_t stream, + PolicySelector policy_selector = {}) +{ + using local_offset_t = ::cuda::std::int32_t; + using global_offset_t = OffsetT; + static constexpr bool use_streaming_invocation = + ::cuda::std::numeric_limits::max() > ::cuda::std::numeric_limits::max(); + using streaming_context_t = ::cuda::std:: + conditional_t, NullType>; + using ScanTileStateT = ReduceByKeyScanTileState; + static constexpr int init_kernel_threads = 128; + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(ptx_compute_cap(cc))) + { + return error; + } + + const RleNonTrivialRunsPolicy active_policy = policy_selector(cc); +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceRle to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const int threads_per_block = active_policy.lookback.threads_per_block; + const int items_per_thread = active_policy.lookback.items_per_thread; + const auto tile_size = + static_cast(threads_per_block) * static_cast(items_per_thread); + + auto capped_num_items_per_invocation = num_items; + if constexpr (use_streaming_invocation) + { + capped_num_items_per_invocation = static_cast(::cuda::std::numeric_limits::max()); + capped_num_items_per_invocation -= (capped_num_items_per_invocation % tile_size); + } + + const auto max_num_items_per_invocation = + use_streaming_invocation ? ::cuda::std::min(capped_num_items_per_invocation, num_items) : num_items; + const auto num_partitions = + (capped_num_items_per_invocation == 0) + ? global_offset_t{1} + : ::cuda::ceil_div(num_items, capped_num_items_per_invocation); + + const int max_num_tiles = static_cast(::cuda::ceil_div(max_num_items_per_invocation, tile_size)); + + size_t allocation_sizes[3]; + if (const auto error = CubDebug(ScanTileStateT::AllocationSize(max_num_tiles, allocation_sizes[0]))) + { + return error; + } + allocation_sizes[1] = num_partitions > 1 ? sizeof(global_offset_t) * 2 : size_t{0}; + allocation_sizes[2] = num_partitions > 1 ? sizeof(length_t) * 2 : size_t{0}; + + void* allocations[3] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + for (global_offset_t partition_idx = 0; partition_idx < num_partitions; partition_idx++) + { + global_offset_t current_partition_offset = partition_idx * capped_num_items_per_invocation; + global_offset_t current_num_items = + (partition_idx + 1 == num_partitions) ? (num_items - current_partition_offset) : capped_num_items_per_invocation; + + const auto num_current_tiles = static_cast(::cuda::ceil_div(current_num_items, tile_size)); + ScanTileStateT tile_status; + if (const auto error = CubDebug(tile_status.Init(num_current_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + const int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(num_current_tiles, init_kernel_threads)); +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking device_scan_init_kernel<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + init_kernel_threads, + (long long) stream); +#endif + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_kernel_threads, 0, stream) + .doit(&detail::scan::DeviceCompactInitKernel, + tile_status, + num_current_tiles, + d_num_runs_out))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + if (num_items <= 1) + { + return cudaSuccess; + } +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking device_rle_sweep_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread\n", + num_current_tiles, + threads_per_block, + (long long) stream, + items_per_thread); +#endif + + auto streaming_context = [&] { + if constexpr (use_streaming_invocation) + { + auto tmp_num_uniques = static_cast(allocations[1]); + auto tmp_prefix = static_cast(allocations[2]); + const bool is_first_partition = (partition_idx == 0); + const bool is_last_partition = (partition_idx + 1 == num_partitions); + const int buffer_selector = partition_idx % 2; + return streaming_context_t{ + is_first_partition, + is_last_partition, + current_partition_offset, + &tmp_prefix[buffer_selector], + &tmp_prefix[buffer_selector ^ 0x01], + &tmp_num_uniques[buffer_selector], + &tmp_num_uniques[buffer_selector ^ 0x01]}; + } + else + { + return NullType{}; + } + }(); + + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_current_tiles, threads_per_block, 0, stream) + .doit(&detail::rle::DeviceRleSweepKernel< + PolicySelector, + InputIteratorT, + OffsetsOutputIteratorT, + LengthsOutputIteratorT, + NumRunsOutputIteratorT, + ScanTileStateT, + EqualityOpT, + local_offset_t, + global_offset_t, + streaming_context_t>, + d_in + current_partition_offset, + d_offsets_out, + d_lengths_out, + d_num_runs_out, + tile_status, + equality_op, + static_cast(current_num_items), + num_current_tiles, + streaming_context))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + return cudaSuccess; +} +} // namespace detail::rle +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_scan.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_scan.cuh new file mode 100644 index 00000000..bb3762ef --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_scan.cuh @@ -0,0 +1,1469 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * @brief cub::DeviceScan provides device-wide, parallel operations for + * computing a prefix scan across a sequence of data items residing + * within device-accessible memory. + */ + +#pragma once + +#include + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::scan +{ +template +struct DeviceScanKernelSource +{ + using ScanTileStateT = ScanTileState; + + CUB_DEFINE_KERNEL_GETTER( + InitKernel, + DeviceScanInitKernel) + + CUB_DEFINE_KERNEL_GETTER( + ScanKernel, + DeviceScanKernel) + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t InputSize() + { + return sizeof(it_value_t); + } + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t InputAlign() + { + return alignof(it_value_t); + } + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t OutputSize() + { + return sizeof(it_value_t); + } + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t OutputAlign() + { + return alignof(it_value_t); + } + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t AccumSize() + { + return sizeof(AccumT); + } + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t AccumAlign() + { + return alignof(AccumT); + } + + CUB_RUNTIME_FUNCTION static ScanTileStateT TileState() + { + return {}; + } + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t lookahead_tile_state_size() + { + return sizeof(warpspeed::tile_state_t); + } + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t lookahead_tile_state_alignment() + { + return alignof(warpspeed::tile_state_t); + } + + CUB_RUNTIME_FUNCTION static constexpr auto make_tile_state_kernel_arg(ScanTileStateT ts) + { + tile_state_kernel_arg_t arg; + ::cuda::std::__construct_at(&arg.lookback, ::cuda::std::move(ts)); + return arg; + } + + CUB_RUNTIME_FUNCTION static constexpr auto + lookahead_make_tile_state_kernel_arg(void* ts, ::cuda::std::uint32_t* atomic_counter = nullptr) + { + tile_state_kernel_arg_t arg; + ::cuda::std::__construct_at( + &arg.lookahead, + lookahead_tile_state_arg_t{static_cast*>(ts), atomic_counter}); + return arg; + } +}; + +// TODO(griwes): remove in CCCL 4.0 when we drop the scan dispatcher after publishing the tuning API +template +_CCCL_HOST_DEVICE_API constexpr auto convert_policy() -> ScanPolicy +{ + // this does not convert any lookahead policy data, which is fine because we merged lookahead scan during the CCCL 3.4 + // development cycle, so it never had user exposure through the policy_hub, and we can just only support it through + // the policy_selector. + using scan_policy_t = typename LegacyActivePolicy::ScanPolicyT; + return ScanPolicy{ + ScanAlgorithm::lookback, + ScanLookbackPolicy{ + scan_policy_t::BLOCK_THREADS, + scan_policy_t::ITEMS_PER_THREAD, + scan_policy_t::LOAD_ALGORITHM, + scan_policy_t::LOAD_MODIFIER, + scan_policy_t::STORE_ALGORITHM, + scan_policy_t::SCAN_ALGORITHM, + detail::lookback_delay_policy_from_type}, + ScanLookaheadPolicy{}}; +} + +// TODO(griwes): remove in CCCL 4.0 when we drop the scan dispatcher after publishing the tuning API +template +struct policy_selector_from_hub +{ + [[nodiscard]] _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability /*cc*/) const -> ScanPolicy + { + return convert_policy(); + } +}; +} // namespace detail::scan + +/****************************************************************************** + * Dispatch + ******************************************************************************/ + +/** + * @brief Utility class for dispatching the appropriately-tuned kernels for + * DeviceScan + * + * Deprecated [Since 3.5] + * + * @tparam InputIteratorT + * Random-access input iterator type for reading scan inputs @iterator + * + * @tparam OutputIteratorT + * Random-access output iterator type for writing scan outputs @iterator + * + * @tparam ScanOpT + * Binary scan functor type having member + * `auto operator()(const T &a, const U &b)` + * + * @tparam InitValueT + * The init_value element type for ScanOpT (cub::NullType for inclusive scans) + * + * @tparam OffsetT + * Unsigned integer type for global offsets + * + * @tparam EnforceInclusive + * Enum flag to specify whether to enforce inclusive scan. + * + */ +// TODO(griwes): Remove in CCCL 4.0 +template < + typename InputIteratorT, + typename OutputIteratorT, + typename ScanOpT, + typename InitValueT, + typename OffsetT, + typename AccumT = ::cuda::std::__accumulator_t, + ::cuda::std::_If<::cuda::std::is_same_v, + cub::detail::it_value_t, + typename InitValueT::value_type>>, + ForceInclusive EnforceInclusive = ForceInclusive::No, + typename PolicyHub = detail::scan:: + policy_hub, detail::it_value_t, AccumT, OffsetT, ScanOpT>, + typename KernelSource = detail::scan::DeviceScanKernelSource< + detail::scan::policy_selector_from_hub, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t, + ScanOpT, + InitValueT, + OffsetT, + AccumT, + EnforceInclusive>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") DispatchScan +{ + static_assert(::cuda::std::is_unsigned_v && sizeof(OffsetT) >= 4, + "DispatchScan only supports unsigned offset types of at least 4-bytes"); + + //--------------------------------------------------------------------- + // Constants and Types + //--------------------------------------------------------------------- + + static constexpr int INIT_KERNEL_THREADS = 128; + + /// Device-accessible allocation of temporary storage. When nullptr, the + /// required allocation size is written to \p temp_storage_bytes and no work + /// is done. + void* d_temp_storage; + + /// Reference to size in bytes of \p d_temp_storage allocation + size_t& temp_storage_bytes; + + /// Iterator to the input sequence of data items + InputIteratorT d_in; + + /// Iterator to the output sequence of data items + OutputIteratorT d_out; + + /// Binary scan functor + ScanOpT scan_op; + + /// Initial value to seed the exclusive scan + InitValueT init_value; + + /// Total number of input items (i.e., the length of \p d_in) + OffsetT num_items; + + /// CUDA stream to launch kernels within. Default is stream0. + cudaStream_t stream; + + int ptx_version; + + KernelSource kernel_source; + + KernelLauncherFactory launcher_factory; + + /** + * + * @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_in + * Iterator to the input sequence of data items + * + * @param[out] d_out + * Iterator to the output sequence of data items + * + * @param[in] num_items + * Total number of input items (i.e., the length of `d_in`) + * + * @param[in] scan_op + * Binary scan functor + * + * @param[in] init_value + * Initial value to seed the exclusive scan + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. + * Default is stream0. + * + * @param[in] kernel_source + * Object specifying implementation kernels + * + * @param[in] launcher_factory + * Object to execute implementation kernels on the given stream + */ + // TODO(griwes): Remove in CCCL 4.0 + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchScan( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + OffsetT num_items, + ScanOpT scan_op, + InitValueT init_value, + cudaStream_t stream, + int ptx_version, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_in(d_in) + , d_out(d_out) + , scan_op(scan_op) + , init_value(init_value) + , num_items(num_items) + , stream(stream) + , ptx_version(ptx_version) + , kernel_source(kernel_source) + , launcher_factory(launcher_factory) + {} + + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t + Invoke(InitKernelT init_kernel, ScanKernelT scan_kernel, ActivePolicyT policy = {}) + { + // `LOAD_LDG` makes in-place execution UB and doesn't lead to better + // performance. + policy.CheckLoadModifier(); + + // Number of input tiles + const int tile_size = policy.Scan().ThreadsPerBlock() * policy.Scan().ItemsPerThread(); + const int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + auto tile_state = kernel_source.TileState(); + + // Specify temporary storage allocation requirements + size_t allocation_sizes[1]; + if (const auto error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0]))) + { + return error; // bytes needed for tile status descriptors + } + + // Compute allocation pointers into the single storage blob (or compute + // the necessary size of the blob) + void* allocations[1] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + // Return if the caller is simply requesting the size of the storage allocation, or the problem is empty + if (d_temp_storage == nullptr || num_items == 0) + { + return cudaSuccess; + } + + // Construct the tile status interface + if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + // Log init_kernel configuration + const int init_grid_size = ::cuda::ceil_div(num_tiles, INIT_KERNEL_THREADS); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, INIT_KERNEL_THREADS, (long long) stream); +#endif // CUB_DEBUG_LOG + + // Invoke init_kernel to initialize tile descriptors + if (const auto error = CubDebug( + launcher_factory(init_grid_size, INIT_KERNEL_THREADS, 0, stream, /* dependent_launch */ ptx_version >= 900) + .doit(init_kernel, kernel_source.make_tile_state_kernel_arg(tile_state), num_tiles))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Get SM occupancy for scan_kernel + int scan_sm_occupancy; + if (const auto error = + CubDebug(launcher_factory.MaxSmOccupancy(scan_sm_occupancy, scan_kernel, policy.Scan().ThreadsPerBlock()))) + { + return error; + } + + // Get max x-dimension of grid + int max_dim_x; + if (const auto error = CubDebug(launcher_factory.MaxGridDimX(max_dim_x))) + { + return error; + } + + // Run grids in epochs (in case number of tiles exceeds max x-dimension + const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x); + for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size) + { +// Log scan_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking %d scan_kernel<<<%d, %d, 0, %lld>>>(), %d items " + "per thread, %d SM occupancy\n", + start_tile, + scan_grid_size, + policy.Scan().ThreadsPerBlock(), + (long long) stream, + policy.Scan().ItemsPerThread(), + scan_sm_occupancy); +#endif // CUB_DEBUG_LOG + + // Invoke scan_kernel + if (const auto error = CubDebug( + launcher_factory( + scan_grid_size, policy.Scan().ThreadsPerBlock(), 0, stream, /* dependent_launch */ ptx_version >= 900) + .doit(scan_kernel, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in), + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out), + kernel_source.make_tile_state_kernel_arg(tile_state), + start_tile, + scan_op, + init_value, + num_items, + /* num_stages, unused */ 1))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; + } + + // do check in separate function, so error message contains the required SMEM in the error novel + template + CUB_RUNTIME_FUNCTION static void __check_smem() + { + static_assert(SMemSizeForSingleStage <= detail::max_smem_per_block, + "Single-stage lookahead scan exceeds architecture independent SMEM (48KiB)"); + } + + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t __invoke_lookahead_algorithm(PolicyGetter policy_getter) + { +#if __cccl_ptx_isa >= 860 + if (num_items == 0) + { + temp_storage_bytes = 1; // just fulfill the contract that CUB always requires some temporary storage + return cudaSuccess; + } + + CUB_DETAIL_CONSTEXPR_ISH const ScanLookaheadPolicy lookahead_policy = policy_getter().lookahead; + CUB_DETAIL_STATIC_ISH_ASSERT(lookahead_policy.reduce_and_scan_warps >= 1, + "Lookahead scan policy have at least 1 warp for reducing and scanning"); + CUB_DETAIL_STATIC_ISH_ASSERT( + lookahead_policy.items_per_thread >= 1, "Lookahead scan policy must have at least 1 item per thread"); + CUB_DETAIL_STATIC_ISH_ASSERT(lookahead_policy.lookahead_items_per_thread >= 1, + "Lookahead scan policy must look ahead at least 1 item per thread"); + + const int grid_dim = + static_cast(::cuda::ceil_div(num_items, static_cast(lookahead_policy.tile_size()))); + + if (d_temp_storage == nullptr) + { + temp_storage_bytes = static_cast(grid_dim) * kernel_source.lookahead_tile_state_size(); + return cudaSuccess; + } + + if (num_items == 0) + { + return cudaSuccess; + } + + int sm_count = 0; + if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) + { + return error; + } + // Maximum dynamic shared memory size that we can use for temporary storage. + int max_dynamic_smem_size{}; + if (const auto error = + CubDebug(launcher_factory.max_dynamic_smem_size_for(max_dynamic_smem_size, kernel_source.ScanKernel()))) + { + return error; + } + + // TODO(bgruber): we probably need to ensure alignment of d_temp_storage + _CCCL_ASSERT(::cuda::is_aligned(d_temp_storage, kernel_source.lookahead_tile_state_alignment()), ""); + + auto scan_kernel = kernel_source.ScanKernel(); + [[maybe_unused]] auto kernel_src = kernel_source; // need to pull a copy to not access `this` during const. eval. + CUB_DETAIL_CONSTEXPR_ISH int smem_size_1_stage = detail::scan::smem_for_stages( + lookahead_policy, + 1, + static_cast(kernel_src.InputSize()), + static_cast(kernel_src.InputAlign()), + static_cast(kernel_src.OutputAlign()), + static_cast(kernel_src.AccumSize()), + static_cast(kernel_src.AccumAlign())); +# if defined(CUB_DEFINE_RUNTIME_POLICIES) + _CCCL_ASSERT(smem_size_1_stage <= int{detail::max_smem_per_block}, + "Single-stage lookahead scan exceeds architecture independent SMEM (48KiB)"); +# else // defined(CUB_DEFINE_RUNTIME_POLICIES) + __check_smem(); +# endif // defined(CUB_DEFINE_RUNTIME_POLICIES) + + int num_stages = 1; + int smem_size = smem_size_1_stage; + + // When launched from the host, maximize the number of stages that we can fit inside the shared memory. + NV_IF_TARGET(NV_IS_HOST, ({ + // number of stages to have an even workload across all SMs (improves small problem sizes), assuming + // 1 CTA per SM +1 since it tends to improve performance + // TODO(bgruber): make the +1 a tuning parameter + const int max_stages_for_even_workload = static_cast( + ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); + + while (num_stages <= max_stages_for_even_workload) + { + const int next_smem_size = detail::scan::smem_for_stages( + lookahead_policy, + num_stages + 1, + static_cast(kernel_source.InputSize()), + static_cast(kernel_source.InputAlign()), + static_cast(kernel_source.OutputAlign()), + static_cast(kernel_source.AccumSize()), + static_cast(kernel_source.AccumAlign())); + if (next_smem_size > max_dynamic_smem_size) + { + // This number of stages failed, so stay at the current settings + break; + } + + smem_size = next_smem_size; + ++num_stages; + } + + if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, smem_size)) + { + return error; + } + })) + + // Invoke init kernel + { + constexpr auto init_kernel_threads = 128; + const auto init_grid_size = ::cuda::ceil_div(grid_dim, init_kernel_threads); + +# ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceScanInitKernel<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + init_kernel_threads, + (long long) stream); +# endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + launcher_factory(init_grid_size, + init_kernel_threads, + 0, + stream, + /* dependent_launch */ ptx_version >= 900) + .doit(kernel_source.InitKernel(), + kernel_source.lookahead_make_tile_state_kernel_arg(d_temp_storage), + grid_dim))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + // Invoke scan kernel + { + const int block_dim = detail::scan::num_total_threads(lookahead_policy); + +# ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceScanKernel<<<%d, %d, %d, %lld>>>()\n", grid_dim, block_dim, smem_size, (long long) stream); +# endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + launcher_factory(grid_dim, block_dim, smem_size, stream, /* dependent_launch */ ptx_version >= 900) + .doit(scan_kernel, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in), + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out), + kernel_source.lookahead_make_tile_state_kernel_arg(d_temp_storage), + /* start_tile, unused */ 0, + ::cuda::std::move(scan_op), + init_value, + num_items, + num_stages))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } +#else // __cccl_ptx_isa >= 860 + static_assert(sizeof(policy_getter) == 0, + "Implementation bug: Tuning policy selected lookahead, but supported PTX ISA is too low"); +#endif // __cccl_ptx_isa >= 860 + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t __invoke_lookback_algorithm(PolicyGetter policy_getter) + { + CUB_DETAIL_CONSTEXPR_ISH const ScanLookbackPolicy active_policy = policy_getter().lookback; + CUB_DETAIL_STATIC_ISH_ASSERT( + active_policy.threads_per_block >= 1, "Lookback scan policy must have at least 1 thread per block"); + CUB_DETAIL_STATIC_ISH_ASSERT( + active_policy.items_per_thread >= 1, "Lookback scan policy must have at least 1 item per thread"); + CUB_DETAIL_STATIC_ISH_ASSERT(active_policy.load_modifier != CacheLoadModifier::LOAD_LDG, + "The memory consistency model does not apply to texture accesses"); + + // Number of input tiles + const int tile_size = active_policy.threads_per_block * active_policy.items_per_thread; + const int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + auto tile_state = kernel_source.TileState(); + + // Specify temporary storage allocation requirements + size_t allocation_sizes[1]; + if (const auto error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0]))) + { + return error; // bytes needed for tile status descriptors + } + + // Compute allocation pointers into the single storage blob (or compute + // the necessary size of the blob) + void* allocations[1] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + // Return if the caller is simply requesting the size of the storage allocation, or the problem is empty + if (d_temp_storage == nullptr || num_items == 0) + { + return cudaSuccess; + } + + // Construct the tile status interface + if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + // Log init_kernel configuration + constexpr int init_kernel_threads = 128; + const int init_grid_size = ::cuda::ceil_div(num_tiles, init_kernel_threads); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, init_kernel_threads, (long long) stream); +#endif // CUB_DEBUG_LOG + + // Invoke init_kernel to initialize tile descriptors + if (const auto error = CubDebug( + launcher_factory(init_grid_size, init_kernel_threads, 0, stream, /* dependent_launch */ ptx_version >= 900) + .doit(kernel_source.InitKernel(), kernel_source.make_tile_state_kernel_arg(tile_state), num_tiles))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Get SM occupancy for scan_kernel + int scan_sm_occupancy; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + scan_sm_occupancy, kernel_source.ScanKernel(), active_policy.threads_per_block))) + { + return error; + } + + // Get max x-dimension of grid + int max_dim_x; + if (const auto error = CubDebug(launcher_factory.MaxGridDimX(max_dim_x))) + { + return error; + } + + // Run grids in epochs (in case number of tiles exceeds max x-dimension + const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x); + for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size) + { +// Log scan_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking %d scan_kernel<<<%d, %d, 0, %lld>>>(), %d items " + "per thread, %d SM occupancy\n", + start_tile, + scan_grid_size, + active_policy.threads_per_block, + (long long) stream, + active_policy.items_per_thread, + scan_sm_occupancy); +#endif // CUB_DEBUG_LOG + + // Invoke scan_kernel + if (const auto error = CubDebug( + launcher_factory( + scan_grid_size, active_policy.threads_per_block, 0, stream, /* dependent_launch */ ptx_version >= 900) + .doit(kernel_source.ScanKernel(), + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in), + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out), + kernel_source.make_tile_state_kernel_arg(tile_state), + start_tile, + scan_op, + init_value, + num_items, + /* num_stages, unused */ 1))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT = {}) + { + struct policy_getter + { + // host-device not api, because clang warns about exclude_from_explicit_instantiation in local types + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr auto operator()() const + { + return detail::scan::convert_policy(); + } + }; + + if CUB_DETAIL_CONSTEXPR_ISH (policy_getter{}().algorithm == ScanAlgorithm::lookahead) + { + return __invoke_lookahead_algorithm(policy_getter{}); + } + else + { + return __invoke_lookback_algorithm(policy_getter{}); + } + } + + /** + * @brief Internal dispatch routine + * + * @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_in + * Iterator to the input sequence of data items + * + * @param[out] d_out + * Iterator to the output sequence of data items + * + * @param[in] scan_op + * Binary scan functor + * + * @param[in] init_value + * Initial value to seed the exclusive scan + * + * @param[in] num_items + * Total number of input items (i.e., the length of `d_in`) + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. + * Default is stream0. + * + * @param[in] kernel_source + * Object specifying implementation kernels + * + * @param[in] launcher_factory + * Object to execute implementation kernels on the given stream + * + * @param[in] max_policy + * Struct encoding chain of algorithm tuning policies + */ + // TODO(griwes): Remove in CCCL 4.0 + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) + { + // Get PTX version + int ptx_version = 0; + if (const auto error = CubDebug(launcher_factory.PtxVersion(ptx_version))) + { + return error; + } + + // Create dispatch functor + DispatchScan dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_items, + scan_op, + init_value, + stream, + ptx_version, + kernel_source, + launcher_factory); + + // Dispatch to chained policy + return CubDebug(max_policy.Invoke(ptx_version, dispatch)); + } +}; + +namespace detail::scan +{ +// do check in separate function, so error message contains the required SMEM in the error novel +template +CUB_RUNTIME_FUNCTION void check_lookahead_smem() +{ + static_assert(SMemSizeForSingleStage <= detail::max_smem_per_block, + "Single-stage lookahead scan exceeds architecture independent SMEM (48KiB)"); +} + +template +CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t invoke_lookback( + PolicyGetter policy_getter, + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + bool dependent_launch, + KernelSource kernel_source, + KernelLauncherFactory launcher_factory) +{ + CUB_DETAIL_CONSTEXPR_ISH const ScanLookbackPolicy active_policy = policy_getter().lookback; + CUB_DETAIL_STATIC_ISH_ASSERT( + active_policy.threads_per_block >= 1, "Lookback scan policy must have at least 1 thread per block"); + CUB_DETAIL_STATIC_ISH_ASSERT( + active_policy.items_per_thread >= 1, "Lookback scan policy must have at least 1 item per thread"); + CUB_DETAIL_STATIC_ISH_ASSERT(active_policy.load_modifier != CacheLoadModifier::LOAD_LDG, + "The memory consistency model does not apply to texture accesses"); + + // Number of input tiles + const int tile_size = active_policy.threads_per_block * active_policy.items_per_thread; + const int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + auto tile_state = kernel_source.TileState(); + + // Specify temporary storage allocation requirements + size_t allocation_sizes[1]; + if (const auto error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0]))) + { + return error; // bytes needed for tile status descriptors + } + + // Compute allocation pointers into the single storage blob (or compute + // the necessary size of the blob) + void* allocations[1] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + // Return if the caller is simply requesting the size of the storage allocation, or the problem is empty + if (d_temp_storage == nullptr || num_items == 0) + { + return cudaSuccess; + } + + // Construct the tile status interface + if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + // Log init_kernel configuration + constexpr int init_kernel_threads = 128; + const int init_grid_size = ::cuda::ceil_div(num_tiles, init_kernel_threads); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, init_kernel_threads, (long long) stream); +#endif // CUB_DEBUG_LOG + + // Invoke init_kernel to initialize tile descriptors + if (const auto error = CubDebug( + launcher_factory(init_grid_size, init_kernel_threads, 0, stream, dependent_launch) + .doit(kernel_source.InitKernel(), kernel_source.make_tile_state_kernel_arg(tile_state), num_tiles))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Get SM occupancy for scan_kernel + int scan_sm_occupancy; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + scan_sm_occupancy, kernel_source.ScanKernel(), active_policy.threads_per_block))) + { + return error; + } + + // Get max x-dimension of grid + int max_dim_x; + if (const auto error = CubDebug(launcher_factory.MaxGridDimX(max_dim_x))) + { + return error; + } + + // Run grids in epochs (in case number of tiles exceeds max x-dimension + const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x); + for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size) + { +// Log scan_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking %d scan_kernel<<<%d, %d, 0, %lld>>>(), %d items " + "per thread, %d SM occupancy\n", + start_tile, + scan_grid_size, + active_policy.threads_per_block, + (long long) stream, + active_policy.items_per_thread, + scan_sm_occupancy); +#endif // CUB_DEBUG_LOG + + // Invoke scan_kernel + if (const auto error = CubDebug( + launcher_factory(scan_grid_size, active_policy.threads_per_block, 0, stream, dependent_launch) + .doit(kernel_source.ScanKernel(), + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in), + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out), + kernel_source.make_tile_state_kernel_arg(tile_state), + start_tile, + scan_op, + init_value, + num_items, + /* num_stages, unused */ 1))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} + +template +CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t invoke_lookahead( + PolicyGetter policy_getter, + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + bool dependent_launch, + bool atomic_scheduling, + KernelSource kernel_source, + KernelLauncherFactory launcher_factory) +{ +#if __cccl_ptx_isa >= 860 + if (num_items == 0) + { + temp_storage_bytes = 1; // just fulfill the contract that CUB always requires some temporary storage + return cudaSuccess; + } + + CUB_DETAIL_CONSTEXPR_ISH const ScanLookaheadPolicy lookahead_policy = policy_getter().lookahead; + CUB_DETAIL_STATIC_ISH_ASSERT(lookahead_policy.reduce_and_scan_warps >= 1, + "Lookahead scan policy have at least 1 warp for reducing and scanning"); + CUB_DETAIL_STATIC_ISH_ASSERT( + lookahead_policy.items_per_thread >= 1, "Lookahead scan policy must have at least 1 item per thread"); + CUB_DETAIL_STATIC_ISH_ASSERT(lookahead_policy.lookahead_items_per_thread >= 1, + "Lookahead scan policy must look ahead at least 1 item per thread"); + + const int num_tiles = + static_cast(::cuda::ceil_div(num_items, static_cast(lookahead_policy.tile_size()))); + + size_t allocation_sizes[2] = { + static_cast(num_tiles) * kernel_source.lookahead_tile_state_size(), sizeof(::cuda::std::uint32_t)}; + void* allocations[2] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + void* d_tile_state = allocations[0]; + ::cuda::std::uint32_t* d_atomic_counter = static_cast<::cuda::std::uint32_t*>(allocations[1]); + + int sm_count = 0; + if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count))) + { + return error; + } + + const int scan_grid_dim = atomic_scheduling ? (::cuda::std::min) (sm_count, num_tiles) : num_tiles; + // Maximum dynamic shared memory size that we can use for temporary storage. + int max_dynamic_smem_size{}; + if (const auto error = + CubDebug(launcher_factory.max_dynamic_smem_size_for(max_dynamic_smem_size, kernel_source.ScanKernel()))) + { + return error; + } + + auto scan_kernel = kernel_source.ScanKernel(); + [[maybe_unused]] auto kernel_src = kernel_source; // need to pull a copy to not access `this` during const. eval. + CUB_DETAIL_CONSTEXPR_ISH int smem_size_1_stage = detail::scan::smem_for_stages( + lookahead_policy, + 1, + static_cast(kernel_src.InputSize()), + static_cast(kernel_src.InputAlign()), + static_cast(kernel_src.OutputAlign()), + static_cast(kernel_src.AccumSize()), + static_cast(kernel_src.AccumAlign())); +# if defined(CUB_DEFINE_RUNTIME_POLICIES) + _CCCL_ASSERT(smem_size_1_stage <= int{detail::max_smem_per_block}, + "Single-stage lookahead scan exceeds architecture independent SMEM (48KiB)"); +# else // defined(CUB_DEFINE_RUNTIME_POLICIES) + check_lookahead_smem(); +# endif // defined(CUB_DEFINE_RUNTIME_POLICIES) + + int num_stages = 1; + int smem_size = smem_size_1_stage; + + // When launched from the host, maximize the number of stages that we can fit inside the shared memory. + NV_IF_TARGET(NV_IS_HOST, ({ + // number of stages to have an even workload across all SMs (improves small problem sizes), assuming + // 1 CTA per SM +1 since it tends to improve performance + // TODO(bgruber): make the +1 a tuning parameter + const int max_stages_for_even_workload = static_cast( + ::cuda::ceil_div(num_items, static_cast(sm_count * lookahead_policy.tile_size())) + 1); + + while (num_stages <= max_stages_for_even_workload) + { + const int next_smem_size = detail::scan::smem_for_stages( + lookahead_policy, + num_stages + 1, + static_cast(kernel_source.InputSize()), + static_cast(kernel_source.InputAlign()), + static_cast(kernel_source.OutputAlign()), + static_cast(kernel_source.AccumSize()), + static_cast(kernel_source.AccumAlign())); + if (next_smem_size > max_dynamic_smem_size) + { + // This number of stages failed, so stay at the current settings + break; + } + + smem_size = next_smem_size; + ++num_stages; + } + + if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, smem_size)) + { + return error; + } + })) + + // Invoke init kernel + { + constexpr auto init_kernel_threads = 128; + const auto init_grid_size = ::cuda::ceil_div(num_tiles, init_kernel_threads); + +# ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceScanInitKernel<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + init_kernel_threads, + (long long) stream); +# endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + launcher_factory(init_grid_size, init_kernel_threads, 0, stream, dependent_launch) + .doit(kernel_source.InitKernel(), + kernel_source.lookahead_make_tile_state_kernel_arg(d_tile_state, d_atomic_counter), + num_tiles))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + // Invoke scan kernel + { + const int block_dim = detail::scan::num_total_threads(lookahead_policy); +# ifdef CUB_DEBUG_LOG + _CubLog( + "Invoking DeviceScanKernel<<<%d, %d, %d, %lld>>>()\n", scan_grid_dim, block_dim, smem_size, (long long) stream); +# endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + launcher_factory(scan_grid_dim, block_dim, smem_size, stream, dependent_launch) + .doit(scan_kernel, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in), + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out), + kernel_source.lookahead_make_tile_state_kernel_arg(d_tile_state, d_atomic_counter), + /* start_tile, unused */ 0, + ::cuda::std::move(scan_op), + init_value, + num_items, + num_stages))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } +#else // __cccl_ptx_isa >= 860 + static_assert(sizeof(policy_getter) == 0, + "Implementation bug: Tuning policy selected lookahead, but supported PTX ISA is too low"); +#endif // __cccl_ptx_isa >= 860 + return cudaSuccess; +} + +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke( + PolicyGetter policy_getter, + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + ::cuda::compute_capability cc, + KernelSource kernel_source, + KernelLauncherFactory launcher_factory) +{ + const bool dependent_launch = cc >= ::cuda::compute_capability{9, 0}; + if CUB_DETAIL_CONSTEXPR_ISH (policy_getter().algorithm == ScanAlgorithm::lookahead) + { + const bool atomic_scheduling = cc == ::cuda::compute_capability{9, 0}; + return invoke_lookahead( + policy_getter, + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + init_value, + num_items, + stream, + dependent_launch, + atomic_scheduling, + kernel_source, + launcher_factory); + } + else + { + return invoke_lookback( + policy_getter, + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + init_value, + num_items, + stream, + dependent_launch, + kernel_source, + launcher_factory); + } +} + +template < + ForceInclusive EnforceInclusive = ForceInclusive::No, + bool StableReductionOrder = false, + typename InputIteratorT, + typename OutputIteratorT, + typename ScanOpT, + typename InitValueT, + typename OffsetT, + typename AccumT = ::cuda::std::__accumulator_t, + ::cuda::std::_If<::cuda::std::is_same_v, + cub::detail::it_value_t, + typename InitValueT::value_type>>, + typename PolicySelector = + policy_selector_from_types, + typename KernelSource = DeviceScanKernelSource< + PolicySelector, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t, + ScanOpT, + InitValueT, + OffsetT, + AccumT, + EnforceInclusive, + StableReductionOrder>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires scan_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) -> cudaError_t +{ + static_assert(::cuda::std::is_unsigned_v && sizeof(OffsetT) >= 4, + "DispatchScan only supports unsigned offset types of at least 4-bytes"); + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << policy_selector(cc); + _CubLog("Dispatching DeviceScan to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + return dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) { + return invoke( + policy_getter, + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + init_value, + num_items, + stream, + cc, + kernel_source, + launcher_factory); + }); +} + +template , + typename KernelSource = DeviceScanKernelSource< + PolicySelector, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t, + THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t, + ScanOpT, + InitValueT, + OffsetT, + AccumT, + EnforceInclusive, + StableReductionOrder>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch_with_accum( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) -> cudaError_t +{ + return dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + scan_op, + init_value, + num_items, + stream, + policy_selector, + kernel_source, + launcher_factory); +} +} // namespace detail::scan + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_scan_by_key.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_scan_by_key.cuh new file mode 100644 index 00000000..bb316143 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_scan_by_key.cuh @@ -0,0 +1,905 @@ +// SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * @brief DeviceScan provides device-wide, parallel operations for computing a + * prefix scan across a sequence of data items residing within + * device-accessible memory. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +_CCCL_DIAG_PUSH +_CCCL_DIAG_SUPPRESS_GCC("-Wattributes") // __visibility__ attribute ignored +_CCCL_DIAG_SUPPRESS_NVHPC(attribute_requires_external_linkage) + +CUB_NAMESPACE_BEGIN + +/****************************************************************************** + * Kernel entry points + *****************************************************************************/ + +namespace detail::scan_by_key +{ +/** + * @brief Scan by key kernel entry point (multi-block) + * + * @tparam PolicySelector + * Policy selector type + * + * @tparam KeysInputIteratorT + * Random-access input iterator type + * + * @tparam ValuesInputIteratorT + * Random-access input iterator type + * + * @tparam ValuesOutputIteratorT + * Random-access output iterator type + * + * @tparam ScanByKeyTileStateT + * Tile status interface type + * + * @tparam EqualityOp + * Equality functor type + * + * @tparam ScanOpT + * Scan functor type + * + * @tparam InitValueT + * The init_value element for ScanOpT type (cub::NullType for inclusive scan) + * + * @tparam OffsetT + * Unsigned integer type for global offsets + * + * @param d_keys_in + * Input keys data + * + * @param d_keys_prev_in + * Predecessor items for each tile + * + * @param d_values_in + * Input values data + * + * @param d_values_out + * Output values data + * + * @param tile_state + * Tile status interface + * + * @param start_tile + * The starting tile for the current grid + * + * @param equality_op + * Binary equality functor + * + * @param scan_op + * Binary scan functor + * + * @param init_value + * Initial value to seed the exclusive scan + * + * @param num_items + * Total number of scan items for the entire problem + */ +template > +__launch_bounds__(int(current_policy().lookback.threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceScanByKeyKernel( + const KeysInputIteratorT d_keys_in, + KeyT* const d_keys_prev_in, + const ValuesInputIteratorT d_values_in, + const ValuesOutputIteratorT d_values_out, + ScanByKeyTileStateT tile_state, + const int start_tile, + EqualityOp equality_op, + const ScanOpT scan_op, + const InitValueT init_value, + const OffsetT num_items) +{ + static constexpr ScanByKeyPolicy policy = current_policy(); + + using scan_by_key_policy_t = agent_scan_by_key_policy< + policy.lookback.threads_per_block, + policy.lookback.items_per_thread, + policy.lookback.load_algorithm, + policy.lookback.load_modifier, + policy.lookback.scan_algorithm, + policy.lookback.store_algorithm, + delay_constructor_t>; + + // Thread block type for scanning input tiles + using AgentScanByKeyT = detail::scan_by_key::AgentScanByKey< + scan_by_key_policy_t, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>; + + // Shared memory for AgentScanByKey + __shared__ typename AgentScanByKeyT::TempStorage temp_storage; + + // Process tiles + AgentScanByKeyT(temp_storage, d_keys_in, d_keys_prev_in, d_values_in, d_values_out, equality_op, scan_op, init_value) + .ConsumeRange(num_items, tile_state, start_tile); +} + +template +_CCCL_KERNEL_ATTRIBUTES void DeviceScanByKeyInitKernel( + ScanTileStateT tile_state, + const KeysInputIteratorT d_keys_in, + cub::detail::it_value_t* d_keys_prev_in, + const OffsetT items_per_tile, + const int num_tiles) +{ + // Initialize tile status + tile_state.InitializeStatus(num_tiles); + + const int tid = static_cast(blockDim.x * blockIdx.x + threadIdx.x); + const OffsetT tile_base = static_cast(tid) * items_per_tile; + if (tid > 0 && tid < num_tiles) + { + d_keys_prev_in[tid] = d_keys_in[tile_base - 1]; + } +} + +template +struct DeviceScanByKeyKernelSource +{ + using ScanByKeyTileStateT = ReduceByKeyScanTileState; + + CUB_DEFINE_KERNEL_GETTER(InitKernel, DeviceScanByKeyInitKernel) + + CUB_DEFINE_KERNEL_GETTER( + ScanKernel, + DeviceScanByKeyKernel) + + CUB_RUNTIME_FUNCTION static ScanByKeyTileStateT TileState() + { + return {}; + } +}; + +template < + typename KeysInputIteratorT, + typename ValuesInputIteratorT, + typename ValuesOutputIteratorT, + typename EqualityOp, + typename ScanOpT, + typename InitValueT, + typename OffsetT, + typename AccumT = ::cuda::std::__accumulator_t< + ScanOpT, + cub::detail::it_value_t, + ::cuda::std:: + _If<::cuda::std::is_same_v, cub::detail::it_value_t, InitValueT>>, + typename PolicyHub = policy_hub, ScanOpT>, + typename PolicySelector = policy_selector_from_hub, + typename KernelSource = DeviceScanByKeyKernelSource< + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct dispatch_scan_by_key +{ + static_assert(::cuda::std::is_unsigned_v && sizeof(OffsetT) >= 4, + "dispatch_scan_by_key only supports unsigned offset types of at least 4-bytes"); + + //--------------------------------------------------------------------- + // Constants and Types + //--------------------------------------------------------------------- + + static constexpr int INIT_KERNEL_THREADS = 128; + + // The input key type + using KeyT = cub::detail::it_value_t; + + // The input value type + using InputT = cub::detail::it_value_t; + + // Tile state used for the decoupled look-back + using ScanByKeyTileStateT = typename KernelSource::ScanByKeyTileStateT; + + /// Device-accessible allocation of temporary storage. When `nullptr`, the + /// required allocation size is written to `temp_storage_bytes` and no work + /// is done. + void* d_temp_storage; + + /// Reference to size in bytes of `d_temp_storage` allocation + size_t& temp_storage_bytes; + + /// Iterator to the input sequence of key items + KeysInputIteratorT d_keys_in; + + /// Iterator to the input sequence of value items + ValuesInputIteratorT d_values_in; + + /// Iterator to the input sequence of value items + ValuesOutputIteratorT d_values_out; + + /// Binary equality functor + EqualityOp equality_op; + + /// Binary scan functor + ScanOpT scan_op; + + /// Initial value to seed the exclusive scan + InitValueT init_value; + + /// Total number of input items (i.e., the length of `d_in`) + OffsetT num_items; + + /// CUDA stream to launch kernels within. + cudaStream_t stream; + int ptx_version; + KernelSource kernel_source; + KernelLauncherFactory launcher_factory; + + /** + * @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 to the input sequence of key items + * + * @param[in] d_values_in + * Iterator to the input sequence of value items + * + * @param[out] d_values_out + * Iterator to the output sequence of value items + * + * @param[in] equality_op + * Binary equality functor + * + * @param[in] scan_op + * Binary scan functor + * + * @param[in] init_value + * Initial value to seed the exclusive scan + * + * @param[in] num_items + * Total number of input items (i.e., the length of `d_in`) + * + * @param[in] stream + * CUDA stream to launch kernels within. + * + * @param[in] kernel_source + * Object specifying implementation kernels + * + * @param[in] launcher_factory + * Object to execute implementation kernels on the given stream + * + * @param[in] max_policy + * Struct encoding chain of algorithm tuning policies + */ + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE dispatch_scan_by_key( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + EqualityOp equality_op, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + int ptx_version, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_keys_in(d_keys_in) + , d_values_in(d_values_in) + , d_values_out(d_values_out) + , equality_op(equality_op) + , scan_op(scan_op) + , init_value(init_value) + , num_items(num_items) + , stream(stream) + , ptx_version(ptx_version) + , kernel_source(kernel_source) + , launcher_factory(launcher_factory) + {} + + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t __invoke(ScanByKeyPolicy active_policy) + { + // Get device ordinal + int device_ordinal; + if (const auto error = CubDebug(cudaGetDevice(&device_ordinal))) + { + return error; + } + + // Number of input tiles + const int tile_size = active_policy.lookback.threads_per_block * active_policy.lookback.items_per_thread; + const int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + auto tile_state = kernel_source.TileState(); + + // Specify temporary storage allocation requirements + size_t allocation_sizes[2]; + if (const auto error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0]))) + { + return error; // bytes needed for tile status descriptors + } + + allocation_sizes[1] = sizeof(KeyT) * (num_tiles + 1); + + // Compute allocation pointers into the single storage blob (or compute + // the necessary size of the blob) + void* allocations[2] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + // Return if the caller is simply requesting the size of the storage allocation, or the problem is empty + if (d_temp_storage == nullptr || num_items == 0) + { + return cudaSuccess; + } + + KeyT* d_keys_prev_in = static_cast(allocations[1]); + + // Construct the tile status interface + if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + // Log init_kernel configuration + const int init_grid_size = ::cuda::ceil_div(num_tiles, INIT_KERNEL_THREADS); +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, INIT_KERNEL_THREADS, (long long) stream); +#endif // CUB_DEBUG_LOG + + // Invoke init_kernel to initialize tile descriptors + if (const auto error = CubDebug( + launcher_factory(init_grid_size, INIT_KERNEL_THREADS, 0, stream) + .doit(kernel_source.InitKernel(), + tile_state, + d_keys_in, + d_keys_prev_in, + static_cast(tile_size), + num_tiles))) + { + return error; + } + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Get max x-dimension of grid + int max_dim_x; + if (const auto error = CubDebug(cudaDeviceGetAttribute(&max_dim_x, cudaDevAttrMaxGridDimX, device_ordinal))) + { + return error; + } + + // Run grids in epochs (in case number of tiles exceeds max x-dimension + const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x); + for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size) + { + // Log scan_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking %d scan_kernel<<<%d, %d, 0, %lld>>>(), %d items " + "per thread\n", + start_tile, + scan_grid_size, + active_policy.lookback.threads_per_block, + (long long) stream, + active_policy.lookback.items_per_thread); +#endif // CUB_DEBUG_LOG + + // Invoke scan_kernel + if (const auto error = CubDebug( + launcher_factory(scan_grid_size, active_policy.lookback.threads_per_block, 0, stream) + .doit(kernel_source.ScanKernel(), + d_keys_in, + d_keys_prev_in, + d_values_in, + d_values_out, + tile_state, + start_tile, + equality_op, + scan_op, + init_value, + num_items))) + { + return error; + } + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT = {}) + { + return __invoke(detail::scan_by_key::convert_policy()); + } + + /** + * @brief Internal dispatch routine + * + * @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 to the input sequence of key items + * + * @param[in] d_values_in + * Iterator to the input sequence of value items + * + * @param[out] d_values_out + * Iterator to the output sequence of value items + * + * @param[in] equality_op + * Binary equality functor + * + * @param[in] scan_op + * Binary scan functor + * + * @param[in] init_value + * Initial value to seed the exclusive scan + * + * @param[in] num_items + * Total number of input items (i.e., the length of `d_in`) + * + * @param[in] stream + * CUDA stream to launch kernels within. + */ + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + EqualityOp equality_op, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + { + int ptx_version = 0; + if (const auto error = CubDebug(launcher_factory.PtxVersion(ptx_version))) + { + return error; + } + + dispatch_scan_by_key dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + scan_op, + init_value, + num_items, + stream, + ptx_version, + kernel_source, + launcher_factory); + + return CubDebug(typename PolicyHub::MaxPolicy{}.Invoke(ptx_version, dispatch)); + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + EqualityOp equality_op, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + PolicySelectorT policy_selector, + KernelSourceT kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + { + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << policy_selector(cc); + _CubLog("Dispatching DeviceScanByKey to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const ScanByKeyPolicy active_policy = policy_selector(cc); + + return dispatch_scan_by_key< + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT, + PolicyHub, + PolicySelectorT, + KernelSourceT, + KernelLauncherFactory>( + d_temp_storage, + temp_storage_bytes, + d_keys_in, + d_values_in, + d_values_out, + equality_op, + scan_op, + init_value, + num_items, + stream, + -1, + kernel_source, + launcher_factory) + .__invoke(active_policy); + } +}; + +template < + typename OverrideAccumT = use_default, + typename KeysInputIteratorT, + typename ValuesInputIteratorT, + typename ValuesOutputIteratorT, + typename EqualityOp, + typename ScanOpT, + typename InitValueT, + typename OffsetT, + typename AccumT = ::cuda::std::conditional_t< + !::cuda::std::is_same_v, + OverrideAccumT, + ::cuda::std::__accumulator_t< + ScanOpT, + cub::detail::it_value_t, + ::cuda::std:: + _If<::cuda::std::is_same_v, cub::detail::it_value_t, InitValueT>>>, + typename PolicySelector = policy_selector_from_types, + AccumT, + cub::detail::it_value_t, + ScanOpT>, + typename KernelSource = DeviceScanByKeyKernelSource< + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires scan_by_key_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + ValuesInputIteratorT d_values_in, + ValuesOutputIteratorT d_values_out, + EqualityOp equality_op, + ScanOpT scan_op, + InitValueT init_value, + OffsetT num_items, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) -> cudaError_t +{ + static_assert(::cuda::std::is_unsigned_v && sizeof(OffsetT) >= 4, + "scan_by_key::dispatch only supports unsigned offset types of at least 4-bytes"); + + using KeyT = cub::detail::it_value_t; + + static constexpr int INIT_KERNEL_THREADS = 128; + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << policy_selector(cc); + _CubLog("Dispatching DeviceScanByKey to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const ScanByKeyPolicy active_policy = policy_selector(cc); + + // Get device ordinal + int device_ordinal; + if (const auto error = CubDebug(cudaGetDevice(&device_ordinal))) + { + return error; + } + + // Number of input tiles + const int tile_size = active_policy.lookback.threads_per_block * active_policy.lookback.items_per_thread; + const int num_tiles = static_cast(::cuda::ceil_div(num_items, tile_size)); + + auto tile_state = kernel_source.TileState(); + + // Specify temporary storage allocation requirements + size_t allocation_sizes[2]; + if (const auto error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0]))) + { + return error; // bytes needed for tile status descriptors + } + + allocation_sizes[1] = sizeof(KeyT) * (num_tiles + 1); + + // Compute allocation pointers into the single storage blob (or compute + // the necessary size of the blob) + void* allocations[2] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + // Return if the caller is simply requesting the size of the storage allocation, or the problem is empty + if (d_temp_storage == nullptr || num_items == 0) + { + return cudaSuccess; + } + + KeyT* d_keys_prev_in = static_cast(allocations[1]); + + // Construct the tile status interface + if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + // Log init_kernel configuration + const int init_grid_size = ::cuda::ceil_div(num_tiles, INIT_KERNEL_THREADS); +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, INIT_KERNEL_THREADS, (long long) stream); +#endif // CUB_DEBUG_LOG + + // Invoke init_kernel to initialize tile descriptors + if (const auto error = CubDebug( + launcher_factory(init_grid_size, INIT_KERNEL_THREADS, 0, stream) + .doit(kernel_source.InitKernel(), + tile_state, + d_keys_in, + d_keys_prev_in, + static_cast(tile_size), + num_tiles))) + { + return error; + } + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Get max x-dimension of grid + int max_dim_x; + if (const auto error = CubDebug(cudaDeviceGetAttribute(&max_dim_x, cudaDevAttrMaxGridDimX, device_ordinal))) + { + return error; + } + + // Run grids in epochs (in case number of tiles exceeds max x-dimension + const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x); + for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size) + { + // Log scan_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking %d scan_kernel<<<%d, %d, 0, %lld>>>(), %d items " + "per thread\n", + start_tile, + scan_grid_size, + active_policy.lookback.threads_per_block, + (long long) stream, + active_policy.lookback.items_per_thread); +#endif // CUB_DEBUG_LOG + + // Invoke scan_kernel + if (const auto error = CubDebug( + launcher_factory(scan_grid_size, active_policy.lookback.threads_per_block, 0, stream) + .doit(kernel_source.ScanKernel(), + d_keys_in, + d_keys_prev_in, + d_values_in, + d_values_out, + tile_state, + start_tile, + equality_op, + scan_op, + init_value, + num_items))) + { + return error; + } + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} +} // namespace detail::scan_by_key + +// TODO(griwes): remove in CCCL 4.0 +template < + typename KeysInputIteratorT, + typename ValuesInputIteratorT, + typename ValuesOutputIteratorT, + typename EqualityOp, + typename ScanOpT, + typename InitValueT, + typename OffsetT, + typename AccumT = ::cuda::std::__accumulator_t< + ScanOpT, + cub::detail::it_value_t, + ::cuda::std:: + _If<::cuda::std::is_same_v, cub::detail::it_value_t, InitValueT>>, + typename PolicyHub = + detail::scan_by_key::policy_hub, ScanOpT>, + typename PolicySelector = detail::scan_by_key::policy_selector_from_hub, + typename KernelSource = detail::scan_by_key::DeviceScanByKeyKernelSource< + PolicySelector, + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +using DispatchScanByKey + CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") = detail::scan_by_key::dispatch_scan_by_key< + KeysInputIteratorT, + ValuesInputIteratorT, + ValuesOutputIteratorT, + EqualityOp, + ScanOpT, + InitValueT, + OffsetT, + AccumT, + PolicyHub, + PolicySelector, + KernelSource, + KernelLauncherFactory>; + +CUB_NAMESPACE_END + +_CCCL_DIAG_POP diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_radix_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_radix_sort.cuh new file mode 100644 index 00000000..2c9d6fbe --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_radix_sort.cuh @@ -0,0 +1,979 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::DeviceRadixSort provides device-wide, parallel operations for computing a radix sort across + * a sequence of data items residing within device-accessible memory. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// suppress warnings triggered by #pragma unroll: +// "warning: loop not unrolled: the optimizer was unable to perform the requested transformation; the transformation +// might be disabled or specified as part of an unsupported transformation ordering [-Wpass-failed=transform-warning]" +_CCCL_DIAG_PUSH +_CCCL_DIAG_SUPPRESS_CLANG("-Wpass-failed") + +CUB_NAMESPACE_BEGIN + +namespace detail::segmented_radix_sort +{ +template +struct DeviceSegmentedRadixSortKernelSource +{ + static_assert(::cuda::std::is_empty_v); + + CUB_DEFINE_KERNEL_GETTER( + SegmentedRadixSortKernel, + DeviceSegmentedRadixSortKernel< + PolicySelectorT, + false, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + SegmentSizeT, + DecomposerT>); + + CUB_DEFINE_KERNEL_GETTER( + AltSegmentedRadixSortKernel, + DeviceSegmentedRadixSortKernel< + PolicySelectorT, + true, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + SegmentSizeT, + DecomposerT>); + + CUB_RUNTIME_FUNCTION static constexpr size_t KeySize() + { + return sizeof(KeyT); + } + + CUB_RUNTIME_FUNCTION static constexpr size_t ValueSize() + { + return sizeof(ValueT); + } +}; + +// TODO(bgruber): remove in CCCL 4.0 when we drop the radix sort dispatcher after publishing the tuning API +template +_CCCL_HOST_DEVICE_API constexpr auto convert_policy() -> SegmentedRadixSortPolicy +{ + using active_policy = LegacyActivePolicy; + + const auto regular_pass = radix_sort::convert_downsweep_policy(typename active_policy::SegmentedPolicy{}); + const auto alternate_pass = radix_sort::convert_downsweep_policy(typename active_policy::AltSegmentedPolicy{}); + return SegmentedRadixSortPolicy{regular_pass, alternate_pass}; +} + +// TODO(bgruber): remove in CCCL 4.0 when we drop the radix sort dispatcher after publishing the tuning API +template +struct policy_selector_from_hub +{ + _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> SegmentedRadixSortPolicy + { + return convert_policy(); + } +}; +} // namespace detail::segmented_radix_sort + +/****************************************************************************** + * Segmented dispatch + ******************************************************************************/ + +/** + * @brief Utility class for dispatching the appropriately-tuned kernels for segmented device-wide + * radix sort + * + * Deprecated [Since 3.5] + * + * @tparam SortOrder + * Whether to sort in ascending or descending order + * + * @tparam KeyT + * Key type + * + * @tparam ValueT + * Value type + * + * @tparam BeginOffsetIteratorT + * Random-access input iterator type for reading segment beginning offsets @iterator + * + * @tparam EndOffsetIteratorT + * Random-access input iterator type for reading segment ending offsets @iterator + * + * @tparam SegmentSizeT + * Integer type to index items within a segment + */ +// TODO(bgruber): drop in CCCL 4.0 +template , + typename DecomposerT = detail::identity_decomposer_t, + typename KernelSource = detail::segmented_radix_sort::DeviceSegmentedRadixSortKernelSource< + detail::segmented_radix_sort::policy_selector_from_hub, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + SegmentSizeT, + DecomposerT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceSegmentedRadixSort") DispatchSegmentedRadixSort +{ + //------------------------------------------------------------------------------ + // Constants + //------------------------------------------------------------------------------ + + // Whether this is a keys-only (or key-value) sort + static constexpr bool KEYS_ONLY = ::cuda::std::is_same_v; + + //------------------------------------------------------------------------------ + // Parameter members + //------------------------------------------------------------------------------ + + /// Device-accessible allocation of temporary storage. When nullptr, the required allocation size + /// is written to `temp_storage_bytes` and no work is done. + void* d_temp_storage; + + /// Reference to size in bytes of `d_temp_storage` allocation + size_t& temp_storage_bytes; + + /// Double-buffer whose current buffer contains the unsorted input keys and, upon return, is + /// updated to point to the sorted output keys + DoubleBuffer& d_keys; + + /// Double-buffer whose current buffer contains the unsorted input values and, upon return, is + /// updated to point to the sorted output values + DoubleBuffer& d_values; + + /// Number of items to sort + ::cuda::std::int64_t num_items; + + /// The number of segments that comprise the sorting data + ::cuda::std::int64_t num_segments; + + /// Random-access input iterator to the sequence of beginning offsets of length `num_segments`, + /// such that d_begin_offsets[i] is the first element of the ith + /// data segment in d_keys_* and d_values_* + BeginOffsetIteratorT d_begin_offsets; + + /// Random-access input iterator to the sequence of ending offsets of length `num_segments`, + /// such that d_end_offsets[i]-1 is the last element of the ith + /// data segment in d_keys_* and d_values_*. If d_end_offsets[i]-1 + /// <= d_begin_offsets[i], the ith is considered empty. + EndOffsetIteratorT d_end_offsets; + + /// The beginning (least-significant) bit index needed for key comparison + int begin_bit; + + /// The past-the-end (most-significant) bit index needed for key comparison + int end_bit; + + /// CUDA stream to launch kernels within. Default is stream0. + cudaStream_t stream; + + /// PTX version + int ptx_version; + + /// Whether is okay to overwrite source buffers + bool is_overwrite_okay; + + DecomposerT decomposer; + + KernelSource kernel_source; + + KernelLauncherFactory launcher_factory; + + //------------------------------------------------------------------------------ + // Constructors + //------------------------------------------------------------------------------ + + /// Constructor + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchSegmentedRadixSort( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit, + int end_bit, + bool is_overwrite_okay, + cudaStream_t stream, + int ptx_version, + DecomposerT decomposer = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_keys(d_keys) + , d_values(d_values) + , num_items(num_items) + , num_segments(num_segments) + , d_begin_offsets(d_begin_offsets) + , d_end_offsets(d_end_offsets) + , begin_bit(begin_bit) + , end_bit(end_bit) + , stream(stream) + , ptx_version(ptx_version) + , is_overwrite_okay(is_overwrite_okay) + , decomposer(decomposer) + , kernel_source(kernel_source) + , launcher_factory(launcher_factory) + {} + + //------------------------------------------------------------------------------ + // Multi-segment invocation + //------------------------------------------------------------------------------ + + /// Invoke a three-kernel sorting pass at the current bit. + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t InvokePass( + const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + int& current_bit, + PassConfigT& pass_config) + { + cudaError error = cudaSuccess; + + // The number of bits to process in this pass + int pass_bits = ::cuda::std::min(pass_config.radix_bits, (end_bit - current_bit)); + + // The offset type (used to specialize the kernel template), large enough to index any segment within a single + // invocation + using per_invocation_segment_offset_t = ::cuda::std::int32_t; + + // The upper bound of segments that a single kernel invocation will process + constexpr auto max_num_segments_per_invocation = + static_cast<::cuda::std::int64_t>(::cuda::std::numeric_limits::max()); + + // Number of radix sort invocations until all segments have been processed + const auto num_invocations = ::cuda::ceil_div(num_segments, max_num_segments_per_invocation); + + BeginOffsetIteratorT begin_offsets_current_it = d_begin_offsets; + EndOffsetIteratorT end_offsets_current_it = d_end_offsets; + + // Iterate over chunks of segments + for (::cuda::std::int64_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const auto current_segment_offset = invocation_index * max_num_segments_per_invocation; + const auto num_current_segments = + ::cuda::std::min(max_num_segments_per_invocation, num_segments - current_segment_offset); + +// Log kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog( + "Invoking segmented_kernels<<<%lld, %lld, 0, %lld>>>(), " + "%lld items per thread, %lld SM occupancy, " + "current segment offset %lld, current bit %d, bit_grain %d\n", + (long long) num_current_segments, + (long long) pass_config.segmented_config.threads_per_block, + (long long) stream, + (long long) pass_config.segmented_config.items_per_thread, + (long long) pass_config.segmented_config.sm_occupancy, + (long long) current_segment_offset, + current_bit, + pass_bits); +#endif + + launcher_factory( + static_cast(num_current_segments), pass_config.segmented_config.threads_per_block, 0, stream) + .doit(pass_config.segmented_kernel, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + begin_offsets_current_it, + end_offsets_current_it, + current_bit, + pass_bits, + decomposer); + + // Check for failure to launch + error = CubDebug(cudaPeekAtLastError()); + if (cudaSuccess != error) + { + return error; + } + + if (invocation_index + 1 < num_invocations) + { + begin_offsets_current_it += num_current_segments; + end_offsets_current_it += num_current_segments; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(detail::DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + } + + // Update current bit once all segments have been processed for the current pass + current_bit += pass_bits; + + return error; + } + + /// PassConfig data structure + template + struct PassConfig + { + SegmentedKernelT segmented_kernel; + detail::KernelConfig segmented_config; + int radix_bits; + int radix_digits; + + /// Initialize pass configuration + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t InitPassConfig( + SegmentedKernelT segmented_kernel, + int radix_bits, + SegmentedPolicyT policy = {}, + KernelLauncherFactory launcher_factory = {}) + { + this->segmented_kernel = segmented_kernel; + this->radix_bits = radix_bits; + this->radix_digits = 1 << radix_bits; + + return CubDebug(segmented_config.Init(segmented_kernel, policy, launcher_factory)); + } + }; + + /** + * @brief Invocation (run multiple digit passes) + * + * @tparam ActivePolicyT + * Umbrella policy active for the target device + * + * @tparam SegmentedKernelT + * Function type of cub::DeviceSegmentedRadixSortKernel + * + * @param[in] segmented_kernel + * Kernel function pointer to parameterization of cub::DeviceSegmentedRadixSortKernel + * + * @param[in] alt_segmented_kernel + * Alternate kernel function pointer to parameterization of + * cub::DeviceSegmentedRadixSortKernel + */ + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t + InvokePasses(SegmentedKernelT segmented_kernel, SegmentedKernelT alt_segmented_kernel, ActivePolicyT policy = {}) + { + cudaError error = cudaSuccess; + do + { + // Init regular and alternate kernel configurations + PassConfig pass_config, alt_pass_config; + + error = pass_config.InitPassConfig( + segmented_kernel, policy.RadixBits(policy.Segmented()), policy.Segmented(), launcher_factory); + if (error) + { + break; + } + + error = alt_pass_config.InitPassConfig( + alt_segmented_kernel, policy.RadixBits(policy.AltSegmented()), policy.AltSegmented(), launcher_factory); + if (error) + { + break; + } + + // Temporary storage allocation requirements + void* allocations[2] = {}; + size_t allocation_sizes[2] = { + // bytes needed for 3rd keys buffer + (is_overwrite_okay) ? 0 : num_items * kernel_source.KeySize(), + + // bytes needed for 3rd values buffer + (is_overwrite_okay || (KEYS_ONLY)) ? 0 : num_items * sizeof(ValueT), + }; + + // Alias the temporary allocations from the single storage blob (or compute the necessary size of the blob) + error = CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)); + if (cudaSuccess != error) + { + break; + } + + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + if (temp_storage_bytes == 0) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + // Pass planning. Run passes of the alternate digit-size configuration until we have an even multiple of our + // preferred digit size + int radix_bits = policy.RadixBits(policy.Segmented()); + int alt_radix_bits = policy.RadixBits(policy.AltSegmented()); + int num_bits = end_bit - begin_bit; + int num_passes = ::cuda::std::max(::cuda::ceil_div(num_bits, radix_bits), 1); // num_bits may be zero + bool is_num_passes_odd = num_passes & 1; + int max_alt_passes = (num_passes * radix_bits) - num_bits; + int alt_end_bit = ::cuda::std::min(end_bit, begin_bit + (max_alt_passes * alt_radix_bits)); + + DoubleBuffer d_keys_remaining_passes( + (is_overwrite_okay || is_num_passes_odd) ? d_keys.Alternate() : static_cast(allocations[0]), + (is_overwrite_okay) ? d_keys.Current() + : (is_num_passes_odd) ? static_cast(allocations[0]) + : d_keys.Alternate()); + + DoubleBuffer d_values_remaining_passes( + (is_overwrite_okay || is_num_passes_odd) ? d_values.Alternate() : static_cast(allocations[1]), + (is_overwrite_okay) ? d_values.Current() + : (is_num_passes_odd) ? static_cast(allocations[1]) + : d_values.Alternate()); + + // Run first pass, consuming from the input's current buffers + int current_bit = begin_bit; + + error = CubDebug(InvokePass( + d_keys.Current(), + d_keys_remaining_passes.Current(), + d_values.Current(), + d_values_remaining_passes.Current(), + current_bit, + (current_bit < alt_end_bit) ? alt_pass_config : pass_config)); + if (cudaSuccess != error) + { + break; + } + + // Run remaining passes + while (current_bit < end_bit) + { + error = CubDebug(InvokePass( + d_keys_remaining_passes.d_buffers[d_keys_remaining_passes.selector], + d_keys_remaining_passes.d_buffers[d_keys_remaining_passes.selector ^ 1], + d_values_remaining_passes.d_buffers[d_keys_remaining_passes.selector], + d_values_remaining_passes.d_buffers[d_keys_remaining_passes.selector ^ 1], + current_bit, + (current_bit < alt_end_bit) ? alt_pass_config : pass_config)); + if (cudaSuccess != error) + { + break; + } + + // Invert selectors and update current bit + d_keys_remaining_passes.selector ^= 1; + d_values_remaining_passes.selector ^= 1; + } + + // Update selector + if (!is_overwrite_okay) + { + num_passes = 1; // Sorted data always ends up in the other vector + } + + d_keys.selector = (d_keys.selector + num_passes) & 1; + d_values.selector = (d_values.selector + num_passes) & 1; + } while (false); + + return error; + } + + //------------------------------------------------------------------------------ + // Chained policy invocation + //------------------------------------------------------------------------------ + + /// Invocation + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT policy = {}) + { + // Return if empty problem, or if no bits to sort and double-buffering is used + if (num_items == 0 || num_segments == 0 || (begin_bit == end_bit && is_overwrite_okay)) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + // Force kernel code-generation in all compiler passes + return InvokePasses(kernel_source.SegmentedRadixSortKernel(), + kernel_source.AltSegmentedRadixSortKernel(), + detail::radix_sort::MakeRadixSortPolicyWrapper(policy)); + } + + //------------------------------------------------------------------------------ + // Dispatch entrypoints + //------------------------------------------------------------------------------ + + /** + * @brief Internal dispatch routine + * + * @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,out] d_keys + * Double-buffer whose current buffer contains the unsorted input keys and, upon return, is + * updated to point to the sorted output keys + * + * @param[in,out] d_values + * Double-buffer whose current buffer contains the unsorted input values and, upon return, is + * updated to point to the sorted output values + * + * @param[in] num_items + * Number of items to sort + * + * @param[in] num_segments + * The number of segments that comprise the sorting data + * + * @param[in] d_begin_offsets + * Random-access input iterator to the sequence of beginning offsets of length + * `num_segments`, such that d_begin_offsets[i] is the first element of the + * ith data segment in d_keys_* and d_values_* + * + * @param[in] d_end_offsets + * Random-access input iterator to the sequence of ending offsets of length `num_segments`, + * such that d_end_offsets[i]-1 is the last element of the ith + * data segment in d_keys_* and d_values_*. + * If d_end_offsets[i]-1 <= d_begin_offsets[i], + * the ith is considered empty. + * + * @param[in] begin_bit + * The beginning (least-significant) bit index needed for key comparison + * + * @param[in] end_bit + * The past-the-end (most-significant) bit index needed for key comparison + * + * @param[in] is_overwrite_okay + * Whether is okay to overwrite source buffers + * + * @param[in] stream + * CUDA stream to launch kernels within. Default is stream0. + */ + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit, + int end_bit, + bool is_overwrite_okay, + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) + { + cudaError_t error; + do + { + // Get PTX version + int ptx_version = 0; + + error = CubDebug(launcher_factory.PtxVersion(ptx_version)); + if (cudaSuccess != error) + { + break; + } + + // Create dispatch functor + DispatchSegmentedRadixSort dispatch( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + is_overwrite_okay, + stream, + ptx_version, + {}, + kernel_source, + launcher_factory); + + // Dispatch to chained policy + error = CubDebug(max_policy.Invoke(ptx_version, dispatch)); + if (cudaSuccess != error) + { + break; + } + } while (false); + + return error; + } +}; + +namespace detail::segmented_radix_sort +{ +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke_passes( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit, + int end_bit, + bool is_overwrite_okay, + cudaStream_t stream, + DecomposerT decomposer, + SegmentedRadixSortPolicy active_policy, + KernelSource kernel_source, + KernelLauncherFactory launcher_factory) +{ + constexpr bool keys_only = ::cuda::std::is_same_v; + + auto segmented_kernel = kernel_source.SegmentedRadixSortKernel(); + auto alt_segmented_kernel = kernel_source.AltSegmentedRadixSortKernel(); + + KernelConfig seg_config, alt_seg_config; + if (const auto error = CubDebug(seg_config.__init(segmented_kernel, active_policy.regular_pass, launcher_factory))) + { + return error; + } + if (const auto error = + CubDebug(alt_seg_config.__init(alt_segmented_kernel, active_policy.alternate_pass, launcher_factory))) + { + return error; + } + + void* allocations[2] = {}; + size_t allocation_sizes[2] = { + (is_overwrite_okay) ? 0 : num_items * kernel_source.KeySize(), + (is_overwrite_okay || keys_only) ? 0 : num_items * sizeof(ValueT), + }; + + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + if (temp_storage_bytes == 0) // if alias_temporaries computed 0 bytes, make sure we request at least 1 + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + const int radix_bits = active_policy.regular_pass.radix_bits; + const int alt_radix_bits = active_policy.alternate_pass.radix_bits; + const int num_bits = end_bit - begin_bit; + const int num_passes = ::cuda::std::max(::cuda::ceil_div(num_bits, radix_bits), 1); + const bool is_num_passes_odd = num_passes & 1; + const int max_alt_passes = (num_passes * radix_bits) - num_bits; + const int alt_end_bit = ::cuda::std::min(end_bit, begin_bit + (max_alt_passes * alt_radix_bits)); + + DoubleBuffer d_keys_remaining_passes( + (is_overwrite_okay || is_num_passes_odd) ? d_keys.Alternate() : static_cast(allocations[0]), + (is_overwrite_okay) ? d_keys.Current() + : (is_num_passes_odd) ? static_cast(allocations[0]) + : d_keys.Alternate()); + + DoubleBuffer d_values_remaining_passes( + (is_overwrite_okay || is_num_passes_odd) ? d_values.Alternate() : static_cast(allocations[1]), + (is_overwrite_okay) ? d_values.Current() + : (is_num_passes_odd) ? static_cast(allocations[1]) + : d_values.Alternate()); + + // The offset type (used to specialize the kernel template), large enough to index any segment within a single + // invocation + using per_invocation_segment_offset_t = ::cuda::std::int32_t; + + // The upper bound of segments that a single kernel invocation will process + constexpr auto max_num_segments_per_invocation = + static_cast<::cuda::std::int64_t>(::cuda::std::numeric_limits::max()); + + // Number of radix sort invocations until all segments have been processed + const auto num_invocations = ::cuda::ceil_div(num_segments, max_num_segments_per_invocation); + + auto invoke_pass = + [&](const KeyT* d_keys_in, + KeyT* d_keys_out, + const ValueT* d_values_in, + ValueT* d_values_out, + int& current_bit, + int pass_radix_bits, + auto kernel, + KernelConfig config) -> cudaError_t { + // The number of bits to process in this pass + const int pass_bits = ::cuda::std::min(pass_radix_bits, (end_bit - current_bit)); + + BeginOffsetIteratorT begin_it = d_begin_offsets; + EndOffsetIteratorT end_it = d_end_offsets; + + // Iterate over chunks of segments + for (::cuda::std::int64_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const auto current_segment_offset = invocation_index * max_num_segments_per_invocation; + const auto num_current_segments = + ::cuda::std::min(max_num_segments_per_invocation, num_segments - current_segment_offset); + + // Log kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog( + "Invoking segmented_kernels<<<%lld, %lld, 0, %lld>>>(), " + "%lld items per thread, %lld SM occupancy, " + "current segment offset %lld, current bit %d, bit_grain %d\n", + (long long) num_current_segments, + (long long) config.threads_per_block, + (long long) stream, + (long long) config.items_per_thread, + (long long) config.sm_occupancy, + (long long) current_segment_offset, + current_bit, + pass_bits); +#endif + + if (const auto err = CubDebug( + launcher_factory(static_cast(num_current_segments), config.threads_per_block, 0, stream) + .doit(kernel, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + begin_it, + end_it, + current_bit, + pass_bits, + decomposer))) + { + return err; + } + if (const auto err = CubDebug(cudaPeekAtLastError())) + { + return err; + } + + if (invocation_index + 1 < num_invocations) + { + begin_it += num_current_segments; + end_it += num_current_segments; + } + + // Sync the stream if specified to flush runtime errors + if (const auto err = CubDebug(detail::DebugSyncStream(stream))) + { + return err; + } + } + + // Update current bit once all segments have been processed for the current pass + current_bit += pass_bits; + + return cudaSuccess; + }; + + // Run first pass, consuming from the input's current buffers + int current_bit = begin_bit; + const bool use_alt_first = current_bit < alt_end_bit; + if (const auto error = CubDebug(invoke_pass( + d_keys.Current(), + d_keys_remaining_passes.Current(), + d_values.Current(), + d_values_remaining_passes.Current(), + current_bit, + use_alt_first ? alt_radix_bits : radix_bits, + use_alt_first ? alt_segmented_kernel : segmented_kernel, + use_alt_first ? alt_seg_config : seg_config))) + { + return error; + } + + // Run remaining passes + while (current_bit < end_bit) + { + const bool use_alt = current_bit < alt_end_bit; + if (const auto error = CubDebug(invoke_pass( + d_keys_remaining_passes.d_buffers[d_keys_remaining_passes.selector], + d_keys_remaining_passes.d_buffers[d_keys_remaining_passes.selector ^ 1], + d_values_remaining_passes.d_buffers[d_keys_remaining_passes.selector], + d_values_remaining_passes.d_buffers[d_keys_remaining_passes.selector ^ 1], + current_bit, + use_alt ? alt_radix_bits : radix_bits, + use_alt ? alt_segmented_kernel : segmented_kernel, + use_alt ? alt_seg_config : seg_config))) + { + return error; + } + + // Invert selectors and update current bit + d_keys_remaining_passes.selector ^= 1; + d_values_remaining_passes.selector ^= 1; + } + + // Update selector + const int final_num_passes = is_overwrite_okay ? num_passes : 1; + d_keys.selector = (d_keys.selector + final_num_passes) & 1; + d_values.selector = (d_values.selector + final_num_passes) & 1; + + return cudaSuccess; +} + +template , + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + int begin_bit, + int end_bit, + bool is_overwrite_okay, + cudaStream_t stream, + DecomposerT decomposer = {}, + TuningEnvT = {}) +{ + using default_policy_selector_t = policy_selector_from_types; + using policy_selector_t = ::cuda::std::decay_t< + ::cuda::std::execution::__query_result_or_t>; +#if _CCCL_HAS_CONCEPTS() + static_assert(segmented_radix_sort_policy_selector); +#endif // _CCCL_HAS_CONCEPTS() + + auto kernel_source = DeviceSegmentedRadixSortKernelSource< + policy_selector_t, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + SegmentSizeT, + DecomposerT>{}; + auto launcher_factory = KernelLauncherFactory{}; + + if (num_items == 0 || num_segments == 0 || (begin_bit == end_bit && is_overwrite_okay)) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + const SegmentedRadixSortPolicy active_policy = policy_selector_t{}(cc); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceSegmentedRadixSort to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + return invoke_passes( + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + begin_bit, + end_bit, + is_overwrite_okay, + stream, + decomposer, + active_policy, + kernel_source, + launcher_factory); +} +} // namespace detail::segmented_radix_sort + +CUB_NAMESPACE_END + +_CCCL_DIAG_POP diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_reduce.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_reduce.cuh new file mode 100644 index 00000000..35c432b4 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_reduce.cuh @@ -0,0 +1,948 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include // for cub::detail::non_void_value_t, cub::detail::it_value_t + +#include +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::segmented_reduce +{ +template +struct DeviceSegmentedReduceKernelSource +{ + // PolicySelector must be stateless, so we can pass the type to the kernel + static_assert(::cuda::std::is_empty_v); + + CUB_DEFINE_KERNEL_GETTER( + SegmentedReduceKernel, + DeviceSegmentedReduceKernel< + PolicySelector, + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>) +}; + +template +struct policy_selector_from_hub +{ +private: + template + _CCCL_HOST_DEVICE_API static constexpr auto convert_policy() -> SegmentedReducePolicy + { + using rp = typename ActivePolicyT::ReducePolicy; + using srp = typename segmented_reduce::policy_hub::MaxPolicy; + using sp = typename srp::SmallReducePolicy; + using mp = typename srp::MediumReducePolicy; + return SegmentedReducePolicy{ + {rp::BLOCK_THREADS, rp::ITEMS_PER_THREAD, rp::VECTOR_LOAD_LENGTH, rp::BLOCK_ALGORITHM, rp::LOAD_MODIFIER}, + {rp::BLOCK_THREADS, mp::WARP_THREADS, rp::ITEMS_PER_THREAD, rp::VECTOR_LOAD_LENGTH, rp::LOAD_MODIFIER}, + {rp::BLOCK_THREADS, sp::WARP_THREADS, rp::ITEMS_PER_THREAD, rp::VECTOR_LOAD_LENGTH, rp::LOAD_MODIFIER}}; + } + + struct extract_policy_dispatch_t + { + SegmentedReducePolicy& policy; + + template + _CCCL_HOST_DEVICE_API constexpr cudaError_t Invoke() + { + policy = convert_policy(); + return cudaSuccess; + } + }; + +public: + [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr auto operator()(::cuda::compute_capability cc) const + -> SegmentedReducePolicy + { + NV_IF_ELSE_TARGET(NV_IS_HOST, + ({ + const int ptx_version = cc.get() * 10; + SegmentedReducePolicy policy{}; + extract_policy_dispatch_t dispatch{policy}; + PolicyHub::MaxPolicy::Invoke(ptx_version, dispatch); + return policy; + }), + (return convert_policy();)); + } +}; +} // namespace detail::segmented_reduce + +// TODO(bgruber): drop in CCCL 4.0 +/** + * @brief Utility class for dispatching the appropriately-tuned kernels for + * device-wide reduction + * + * Deprecated [Since 3.5] + * + * @tparam InputIteratorT + * Random-access input iterator type for reading input items @iterator + * + * @tparam OutputIteratorT + * Output iterator type for recording the reduced aggregate @iterator + * + * @tparam BeginOffsetIteratorT + * Random-access input iterator type for reading segment beginning offsets + * @iterator + * + * @tparam EndOffsetIteratorT + * Random-access input iterator type for reading segment ending offsets + * @iterator + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam ReductionOpT + * Binary reduction functor type having member + * `auto operator()(const T &a, const U &b)` + * + * @tparam InitValueT + * value type + */ +template < + typename InputIteratorT, + typename OutputIteratorT, + typename BeginOffsetIteratorT, + typename EndOffsetIteratorT, + typename OffsetT, + typename ReductionOpT, + typename InitValueT = cub::detail::non_void_value_t>, + typename AccumT = ::cuda::std::__accumulator_t, InitValueT>, + typename PolicyHub = detail::segmented_reduce::policy_hub, + typename KernelSource = detail::segmented_reduce::DeviceSegmentedReduceKernelSource< + detail::segmented_reduce::policy_selector_from_hub, + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceSegmentedReduce") DispatchSegmentedReduce +{ + //--------------------------------------------------------------------------- + // Problem state + //--------------------------------------------------------------------------- + + /// Device-accessible allocation of temporary storage. When `nullptr`, the + /// required allocation size is written to `temp_storage_bytes` and no work + /// is done. + void* d_temp_storage; + + /// Reference to size in bytes of `d_temp_storage` allocation + size_t& temp_storage_bytes; + + /// Pointer to the input sequence of data items + InputIteratorT d_in; + + /// Pointer to the output aggregate + OutputIteratorT d_out; + + /// The number of segments that comprise the segmented reduction data + ::cuda::std::int64_t num_segments; + + /// Random-access input iterator to the sequence of beginning offsets of + /// length `num_segments`, such that `d_begin_offsets[i]` is the first + /// element of the *i*th data segment in `d_keys_*` and + /// `d_values_*` + BeginOffsetIteratorT d_begin_offsets; + + /// Random-access input iterator to the sequence of ending offsets of length + /// `num_segments`, such that `d_end_offsets[i] - 1` is the last element of + /// the *i*th data segment in `d_keys_*` and `d_values_*`. + /// If `d_end_offsets[i] - 1 <= d_begin_offsets[i]`, the *i*th is + /// considered empty. + EndOffsetIteratorT d_end_offsets; + + /// Binary reduction functor + ReductionOpT reduction_op; + + /// The initial value of the reduction + InitValueT init; + + /// CUDA stream to launch kernels within. Default is stream0. + cudaStream_t stream; + + int ptx_version; + + // Source getter + KernelSource kernel_source; + + KernelLauncherFactory launcher_factory; + + //--------------------------------------------------------------------------- + // Constructor + //--------------------------------------------------------------------------- + + /// Constructor + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchSegmentedReduce( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + int ptx_version, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_in(d_in) + , d_out(d_out) + , num_segments(num_segments) + , d_begin_offsets(d_begin_offsets) + , d_end_offsets(d_end_offsets) + , reduction_op(reduction_op) + , init(init) + , stream(stream) + , ptx_version(ptx_version) + , kernel_source(kernel_source) + , launcher_factory(launcher_factory) + {} + + //--------------------------------------------------------------------------- + // Chained policy invocation + //--------------------------------------------------------------------------- + + /** + * @brief Invocation + * + * @tparam ActivePolicyT + * Umbrella policy active for the target device + * + * @tparam DeviceSegmentedReduceKernelT + * Function type of cub::DeviceSegmentedReduceKernel + * + * @param[in] segmented_reduce_kernel + * Kernel function pointer to instantiation of + * cub::DeviceSegmentedReduceKernel + */ + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t + InvokePasses(DeviceSegmentedReduceKernelT segmented_reduce_kernel, ActivePolicyT policy = {}) + { + cudaError error = cudaSuccess; + + do + { + // Return if the caller is simply requesting the size of the storage + // allocation + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + // Init kernel configuration (computes kernel occupancy) + // maybe only used inside CUB_DEBUG_LOG code sections + [[maybe_unused]] detail::KernelConfig segmented_reduce_config; + error = + CubDebug(segmented_reduce_config.Init(segmented_reduce_kernel, policy.SegmentedReduce(), launcher_factory)); + if (cudaSuccess != error) + { + break; + } + + const auto num_segments_per_invocation = + static_cast<::cuda::std::int64_t>(::cuda::std::numeric_limits<::cuda::std::int32_t>::max()); + const ::cuda::std::int64_t num_invocations = ::cuda::ceil_div(num_segments, num_segments_per_invocation); + + for (::cuda::std::int64_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const auto current_seg_offset = invocation_index * num_segments_per_invocation; + const auto num_current_segments = + ::cuda::std::min(num_segments_per_invocation, num_segments - current_seg_offset); + +// Log device_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking SegmentedDeviceReduceKernel<<<%ld, %d, 0, %lld>>>(), " + "%d items per thread, %d SM occupancy\n", + num_current_segments, + policy.SegmentedReduce().ThreadsPerBlock(), + (long long) stream, + policy.SegmentedReduce().ItemsPerThread(), + segmented_reduce_config.sm_occupancy); +#endif // CUB_DEBUG_LOG + + // Invoke DeviceSegmentedReduceKernel + launcher_factory(static_cast<::cuda::std::uint32_t>(num_current_segments), + policy.SegmentedReduce().ThreadsPerBlock(), + 0, + stream) + .doit(segmented_reduce_kernel, + d_in, + d_out, + d_begin_offsets, + d_end_offsets, + static_cast(num_current_segments), + reduction_op, + init, + 0); + + // Check for failure to launch + error = CubDebug(cudaPeekAtLastError()); + if (cudaSuccess != error) + { + break; + } + + if (invocation_index + 1 < num_invocations) + { + d_out += num_current_segments; + d_begin_offsets += num_current_segments; + d_end_offsets += num_current_segments; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(detail::DebugSyncStream(stream)); + if (cudaSuccess != error) + { + break; + } + } + } while (false); + + return error; + } + + /// Invocation + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT policy = {}) + { + auto wrapped_policy = detail::reduce::MakeReducePolicyWrapper(policy); + // Force kernel code-generation in all compiler passes + return InvokePasses(kernel_source.SegmentedReduceKernel(), wrapped_policy); + } + + //--------------------------------------------------------------------------- + // Dispatch entrypoints + //--------------------------------------------------------------------------- + + /** + * @brief Internal dispatch routine for computing a device-wide reduction + * + * @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_in + * Pointer to the input sequence of data items + * + * @param[out] d_out + * Pointer to the output aggregate + * + * @param[in] num_segments + * The number of segments that comprise the sorting data + * + * @param[in] d_begin_offsets + * Random-access input iterator to the sequence of beginning offsets of + * length `num_segments`, such that `d_begin_offsets[i]` is the first + * element of the *i*th data segment in `d_keys_*` and + * `d_values_*` + * + * @param[in] d_end_offsets + * Random-access input iterator to the sequence of ending offsets of length + * `num_segments`, such that `d_end_offsets[i] - 1` is the last element of + * the *i*th data segment in `d_keys_*` and `d_values_*`. + * If `d_end_offsets[i] - 1 <= d_begin_offsets[i]`, the *i*th is + * considered empty. + * + * @param[in] reduction_op + * Binary reduction functor + * + * @param[in] init + * The initial value of the reduction + * + * @param[in] stream + * **[optional]** CUDA stream to launch kernels within. + * Default is stream0. + */ + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) + { + if (num_segments <= 0) + { + return cudaSuccess; + } + + cudaError error = cudaSuccess; + + do + { + // Get PTX version + int ptx_version = 0; + error = CubDebug(launcher_factory.PtxVersion(ptx_version)); + if (cudaSuccess != error) + { + break; + } + + // Create dispatch functor + DispatchSegmentedReduce dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_out, + num_segments, + d_begin_offsets, + d_end_offsets, + reduction_op, + init, + stream, + ptx_version, + kernel_source, + launcher_factory); + + // Dispatch to chained policy + error = CubDebug(max_policy.Invoke(ptx_version, dispatch)); + if (cudaSuccess != error) + { + break; + } + } while (false); + + return error; + } +}; + +namespace detail::segmented_reduce +{ +// select the accumulator type using an overload set, so __accumulator_t is not instantiated when +// an overriding accumulator type is present. This is needed by CCCL.C. +template +_CCCL_HOST_DEVICE_API auto select_segmented_accum_t(use_default*) + -> ::cuda::std::__accumulator_t, InitValueT>; + +template , int> = 0> +_CCCL_HOST_DEVICE_API auto select_segmented_accum_t(OverrideAccumT*) -> OverrideAccumT; + +template < + typename OverrideAccumT = use_default, + typename OverrideOffsetT = use_default, + typename InputIteratorT, + typename OutputIteratorT, + typename BeginOffsetIteratorT, + typename EndOffsetIteratorT, + // need to evaluate common_iterator_value lazily. This is needed by CCCL.C. + typename OffsetT = typename ::cuda::std::conditional_t<::cuda::std::is_same_v, + common_iterator_value, + ::cuda::std::type_identity>::type, + typename ReductionOpT, + typename InitValueT = non_void_value_t>, + typename AccumT = + decltype(select_segmented_accum_t(static_cast(nullptr))), + typename PolicySelector = policy_selector_from_types, + typename KernelSource = DeviceSegmentedReduceKernelSource< + PolicySelector, + InputIteratorT, + OutputIteratorT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires segmented_reduce_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + ReductionOpT reduction_op, + InitValueT init, + size_t max_segment_size, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + if (num_segments <= 0) + { + return cudaSuccess; + } + + // Get CC + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + const SegmentedReducePolicy active_policy = policy_selector(cc); +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceSegmentedReduce to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + // Compute segments_per_block based on max_segment_size hint + int segments_per_block = 1; + if (max_segment_size != 0) + { + if (::cuda::in_range( + max_segment_size, static_cast(1), static_cast(active_policy.small_reduce.items_per_tile()))) + { + segments_per_block = active_policy.small_reduce.segments_per_block(); + } + else if (::cuda::in_range(max_segment_size, + static_cast(1), + static_cast(active_policy.medium_reduce.items_per_tile()))) + { + segments_per_block = active_policy.medium_reduce.segments_per_block(); + } + } + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + // Init kernel configuration (computes kernel occupancy) + [[maybe_unused]] int sm_occupancy{}; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + sm_occupancy, kernel_source.SegmentedReduceKernel(), active_policy.large_reduce.threads_per_block))) + { + return error; + } + + const auto num_segments_per_invocation = + static_cast<::cuda::std::int64_t>(::cuda::std::numeric_limits<::cuda::std::int32_t>::max()); + const ::cuda::std::int64_t num_invocations = ::cuda::ceil_div(num_segments, num_segments_per_invocation); + + for (::cuda::std::int64_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const auto current_seg_offset = invocation_index * num_segments_per_invocation; + const auto num_current_segments = ::cuda::std::min(num_segments_per_invocation, num_segments - current_seg_offset); + +// Log device_reduce_sweep_kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking SegmentedDeviceReduceKernel<<<%ld, %d, 0, %lld>>>(), " + "%d items per thread, %d SM occupancy\n", + num_current_segments, + active_policy.large_reduce.threads_per_block, + (long long) stream, + active_policy.large_reduce.items_per_thread, + sm_occupancy); +#endif // CUB_DEBUG_LOG + + // Invoke DeviceSegmentedReduceKernel + const auto num_blocks = + ::cuda::ceil_div(num_current_segments, static_cast<::cuda::std::int64_t>(segments_per_block)); + if (const auto error = CubDebug( + launcher_factory( + static_cast<::cuda::std::uint32_t>(num_blocks), active_policy.large_reduce.threads_per_block, 0, stream) + .doit(kernel_source.SegmentedReduceKernel(), + d_in, + d_out, + d_begin_offsets, + d_end_offsets, + static_cast(num_current_segments), + reduction_op, + init, + max_segment_size))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + if (invocation_index + 1 < num_invocations) + { + d_out += num_current_segments; + d_begin_offsets += num_current_segments; + d_end_offsets += num_current_segments; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} + +// @brief Functor to generate a key-value pair from an index and value +template +struct generate_idx_value +{ +private: + Iterator it; + int segment_size; + +public: + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE generate_idx_value(Iterator it, int segment_size) + : it(it) + , segment_size(segment_size) + {} + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE auto operator()(::cuda::std::int64_t idx) const + { + return ::cuda::std::pair(static_cast(idx % segment_size), it[idx]); + } +}; + +template +struct DeviceFixedSizeSegmentedReduceKernelSource +{ + // PolicySelector must be stateless, so we can pass the type to the kernel + static_assert(::cuda::std::is_empty_v); + + CUB_DEFINE_KERNEL_GETTER( + FixedSizeSegmentedReduceKernel, + DeviceFixedSizeSegmentedReduceKernel< + PolicySelector, + InputIteratorT, + OutputIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>) + + CUB_DEFINE_KERNEL_GETTER( + FixedSizeSegmentedReduceKernelFinal, + DeviceFixedSizeSegmentedReduceKernel) + + CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t AccumSize() + { + return sizeof(AccumT); + } +}; + +template >, + typename AccumT = decltype(select_segmented_accum_t( + static_cast(nullptr))), + typename PolicySelector = policy_selector_from_types, + typename KernelSource = DeviceFixedSizeSegmentedReduceKernelSource< + PolicySelector, + InputIteratorT, + OutputIteratorT, + OffsetT, + ReductionOpT, + InitValueT, + AccumT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY, + ::cuda::std::enable_if_t<::cuda::std::is_arithmetic_v, int> = 0> +#if _CCCL_HAS_CONCEPTS() + requires segmented_reduce_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch_fixed_size( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + OutputIteratorT d_out, + ::cuda::std::int64_t num_segments, + OffsetT segment_size, + ReductionOpT reduction_op, + InitValueT init, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + if (num_segments <= 0) + { + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + } + return cudaSuccess; + } + + // Get CC + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + const SegmentedReducePolicy active_policy = policy_selector(cc); +#if !_CCCL_COMPILER(NVRTC) && defined(CUB_DEBUG_LOG) + NV_IF_TARGET( + NV_IS_HOST, + (::std::stringstream ss; ss << active_policy; + _CubLog("Dispatching DeviceFixedSizeSegmentedReduce to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str());)) +#endif // !_CCCL_COMPILER(NVRTC) && defined(CUB_DEBUG_LOG) + + const auto tile_size = active_policy.large_reduce.threads_per_block * active_policy.large_reduce.items_per_thread; + + // Single-phase: segment fits in one tile + if (segment_size < tile_size) + { + int segments_per_block = 1; + + if (segment_size <= active_policy.small_reduce.items_per_tile()) + { + segments_per_block = active_policy.small_reduce.segments_per_block(); + } + else if (segment_size <= active_policy.medium_reduce.items_per_tile()) + { + segments_per_block = active_policy.medium_reduce.segments_per_block(); + } + + if (d_temp_storage == nullptr) + { + temp_storage_bytes = 1; + return cudaSuccess; + } + + constexpr auto num_segments_per_invocation = + static_cast<::cuda::std::int64_t>(::cuda::std::numeric_limits<::cuda::std::int32_t>::max()); + const ::cuda::std::int64_t num_invocations = ::cuda::ceil_div(num_segments, num_segments_per_invocation); + + for (::cuda::std::int64_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const auto current_seg_offset = invocation_index * num_segments_per_invocation; + const auto num_current_segments = + ::cuda::std::min(num_segments_per_invocation, num_segments - current_seg_offset); + const auto num_current_blocks = ::cuda::ceil_div(num_current_segments, segments_per_block); + + if (const auto error = CubDebug( + launcher_factory(static_cast<::cuda::std::int32_t>(num_current_blocks), + active_policy.large_reduce.threads_per_block, + 0, + stream) + .doit(kernel_source.FixedSizeSegmentedReduceKernel(), + d_in, + d_out, + segment_size, + static_cast<::cuda::std::int32_t>(num_current_segments), + reduction_op, + init, + static_cast(nullptr), + 0, + 0))) + { + return error; + } + + d_in += num_segments_per_invocation * segment_size; // NOLINT(bugprone-misplaced-widening-cast) + d_out += num_segments_per_invocation; + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; + } + + // Two-phase: segment spans multiple tiles + const auto tiles_per_segment = static_cast(::cuda::ceil_div(segment_size, tile_size)); + + const auto max_tiles_per_invocation = + static_cast<::cuda::std::int64_t>(::cuda::std::numeric_limits<::cuda::std::int32_t>::max()); + const auto max_segments_per_invocation = max_tiles_per_invocation / tiles_per_segment; + const auto num_invocations = ::cuda::ceil_div(num_segments, max_segments_per_invocation); + const auto num_segments_per_invocation = ::cuda::std::min(max_segments_per_invocation, num_segments); + const auto tiles_per_invocation = num_segments_per_invocation * tiles_per_segment; + + // Temporary storage allocation requirements + void* allocations[1] = {}; + size_t allocation_sizes[1] = {static_cast(tiles_per_invocation) * kernel_source.AccumSize()}; + + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + AccumT* d_block_reductions = static_cast(allocations[0]); + + for (::cuda::std::int64_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const auto current_seg_offset = invocation_index * num_segments_per_invocation; + const auto num_current_segments = ::cuda::std::min(num_segments_per_invocation, num_segments - current_seg_offset); + const auto num_current_blocks = static_cast<::cuda::std::int32_t>(num_current_segments * tiles_per_segment); + + // Phase 1: partial reductions + if (const auto error = CubDebug( + launcher_factory(num_current_blocks, active_policy.large_reduce.threads_per_block, 0, stream) + .doit(kernel_source.FixedSizeSegmentedReduceKernel(), + d_in, + d_out, + segment_size, + num_current_blocks, + reduction_op, + init, + d_block_reductions, + tile_size, + tiles_per_segment))) + { + return error; + } + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Phase 2: final reduction of partial results + const int final_segment_size = tiles_per_segment; + int final_segments_per_block = 1; + + if (final_segment_size <= active_policy.small_reduce.items_per_tile()) + { + final_segments_per_block = active_policy.small_reduce.segments_per_block(); + } + else if (final_segment_size <= active_policy.medium_reduce.items_per_tile()) + { + final_segments_per_block = active_policy.medium_reduce.segments_per_block(); + } + + const auto final_num_current_blocks = ::cuda::ceil_div(num_current_segments, final_segments_per_block); + + if (const auto error = CubDebug( + launcher_factory(static_cast<::cuda::std::int32_t>(final_num_current_blocks), + active_policy.large_reduce.threads_per_block, + 0, + stream) + .doit(kernel_source.FixedSizeSegmentedReduceKernelFinal(), + d_block_reductions, + d_out, + final_segment_size, + static_cast<::cuda::std::int32_t>(num_current_segments), + reduction_op, + init, + static_cast(nullptr), + 0, + 0))) + { + return error; + } + + d_in += num_segments_per_invocation * segment_size; // NOLINT(bugprone-misplaced-widening-cast) + d_out += num_segments_per_invocation; + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} +} // namespace detail::segmented_reduce + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_sort.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_sort.cuh new file mode 100644 index 00000000..409a3bd6 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_segmented_sort.cuh @@ -0,0 +1,1544 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::segmented_sort +{ +// Continuation is called after the partitioning stage. It launches kernels to sort large and small segments using the +// partitioning results. Separation of this stage is required to eliminate device-side synchronization in the CDP mode. +template +CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN cudaError_t device_segmented_sort_continuation( + LargeKernelT large_kernel, + SmallKernelT small_kernel, + int num_segments, + KeyT* d_current_keys, + KeyT* d_final_keys, + device_double_buffer d_keys_double_buffer, + ValueT* d_current_values, + ValueT* d_final_values, + device_double_buffer d_values_double_buffer, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + local_segment_index_t* group_sizes, + local_segment_index_t* large_and_medium_segments_indices, + local_segment_index_t* small_segments_indices, + cudaStream_t stream, + KernelLauncherFactory launcher_factory, + int large_threads_per_block, + int small_threads_per_block, + int medium_segments_per_block, + int small_segments_per_block) +{ + using local_segment_index_t = local_segment_index_t; + const local_segment_index_t large_segments = group_sizes[0]; + + if (large_segments > 0) + { + // One CTA per segment + const local_segment_index_t blocks_in_grid = large_segments; + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking " + "DeviceSegmentedSortKernelLarge<<<%d, %d, 0, %lld>>>()\n", + static_cast(blocks_in_grid), + large_threads_per_block, + (long long) stream); +#endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + launcher_factory(blocks_in_grid, large_threads_per_block, 0, stream) + .doit(large_kernel, + large_and_medium_segments_indices, + d_current_keys, + d_final_keys, + d_keys_double_buffer, + d_current_values, + d_final_values, + d_values_double_buffer, + d_begin_offsets, + d_end_offsets))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(DebugSyncStream(stream))) + { + return error; + } + } + + const local_segment_index_t small_segments = group_sizes[1]; + const local_segment_index_t medium_segments = + static_cast(num_segments) - (large_segments + small_segments); + + const local_segment_index_t small_blocks = ::cuda::ceil_div(small_segments, small_segments_per_block); + + const local_segment_index_t medium_blocks = ::cuda::ceil_div(medium_segments, medium_segments_per_block); + + const local_segment_index_t small_and_medium_blocks_in_grid = small_blocks + medium_blocks; + + if (small_and_medium_blocks_in_grid) + { +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking " + "DeviceSegmentedSortKernelSmall<<<%d, %d, 0, %lld>>>()\n", + static_cast(small_and_medium_blocks_in_grid), + small_threads_per_block, + (long long) stream); +#endif // CUB_DEBUG_LOG + + launcher_factory(small_and_medium_blocks_in_grid, small_threads_per_block, 0, stream) + .doit(small_kernel, + small_segments, + medium_segments, + medium_blocks, + small_segments_indices, + large_and_medium_segments_indices + num_segments - medium_segments, + d_current_keys, + d_final_keys, + d_current_values, + d_final_values, + d_begin_offsets, + d_end_offsets); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} + +#ifdef CUB_RDC_ENABLED +/* + * Continuation kernel is used only in the CDP mode. It's used to + * launch device_segmented_sort_continuation as a separate kernel. + */ +template +__launch_bounds__(1) _CCCL_KERNEL_ATTRIBUTES void DeviceSegmentedSortContinuationKernel( + const LargeKernelT large_kernel, + const SmallKernelT small_kernel, + const local_segment_index_t num_segments, + KeyT* const d_current_keys, + KeyT* const d_final_keys, + device_double_buffer d_keys_double_buffer, + ValueT* const d_current_values, + ValueT* const d_final_values, + device_double_buffer d_values_double_buffer, + const BeginOffsetIteratorT d_begin_offsets, + const EndOffsetIteratorT d_end_offsets, + local_segment_index_t* const group_sizes, + local_segment_index_t* const large_and_medium_segments_indices, + local_segment_index_t* const small_segments_indices, + const KernelLauncherFactory launcher_factory, + const int large_threads_per_block, + const int small_threads_per_block, + const int medium_segments_per_block, + const int small_segments_per_block) +{ + // In case of CDP: + // 1. each CTA has a different main stream + // 2. all streams are non-blocking + // 3. child grid always completes before the parent grid + // 4. streams can be used only from the CTA in which they were created + // 5. streams created on the host cannot be used on the device + // + // Due to (4, 5), we can't pass the user-provided stream in the continuation. + // Due to (1, 2, 3) it's safe to pass the main stream. + [[maybe_unused]] const auto error = CubDebug(detail::segmented_sort::device_segmented_sort_continuation( + large_kernel, + small_kernel, + num_segments, + d_current_keys, + d_final_keys, + d_keys_double_buffer, + d_current_values, + d_final_values, + d_values_double_buffer, + d_begin_offsets, + d_end_offsets, + group_sizes, + large_and_medium_segments_indices, + small_segments_indices, + 0, // always launching on the main stream (see motivation above) + launcher_factory, + large_threads_per_block, + small_threads_per_block, + medium_segments_per_block, + small_segments_per_block)); +} +#endif // CUB_RDC_ENABLED + +template +struct DeviceSegmentedSortKernelSource +{ + CUB_DEFINE_KERNEL_GETTER( + SegmentedSortFallbackKernel, + DeviceSegmentedSortFallbackKernel); + + CUB_DEFINE_KERNEL_GETTER( + SegmentedSortKernelSmall, + DeviceSegmentedSortKernelSmall); + + CUB_DEFINE_KERNEL_GETTER( + SegmentedSortKernelLarge, + DeviceSegmentedSortKernelLarge); + + CUB_RUNTIME_FUNCTION static constexpr size_t KeySize() + { + return sizeof(KeyT); + } + + using LargeSegmentsSelectorT = + cub::detail::segmented_sort::LargeSegmentsSelectorT; + using SmallSegmentsSelectorT = + cub::detail::segmented_sort::SmallSegmentsSelectorT; + + CUB_RUNTIME_FUNCTION static constexpr auto LargeSegmentsSelector( + OffsetT offset, BeginOffsetIteratorT begin_offset_iterator, EndOffsetIteratorT end_offset_iterator) + { + return LargeSegmentsSelectorT(offset, begin_offset_iterator, end_offset_iterator); + } + + CUB_RUNTIME_FUNCTION static constexpr auto SmallSegmentsSelector( + OffsetT offset, BeginOffsetIteratorT begin_offset_iterator, EndOffsetIteratorT end_offset_iterator) + { + return SmallSegmentsSelectorT(offset, begin_offset_iterator, end_offset_iterator); + } + + template + CUB_RUNTIME_FUNCTION static constexpr void + SetSegmentOffset(SelectorT& selector, global_segment_offset_t base_segment_offset) + { + selector.base_segment_offset = base_segment_offset; + } +}; + +// TODO(bgruber): remove in CCCL 4.0 +template +struct policy_selector_from_hub +{ + using max_policy = typename PolicyHub::MaxPolicy; + + // this is only called in device code, so we can ignore the cc parameter + _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability) const -> SegmentedSortPolicy + { + using ap = typename PolicyHub::MaxPolicy::ActivePolicy; + using lp = typename ap::LargeSegmentPolicy; + using sp = typename ap::SmallSegmentPolicy; + using mp = typename ap::MediumSegmentPolicy; + + return SegmentedSortPolicy{ + SegmentedSortRadixSortPolicy{ + lp::BLOCK_THREADS, + lp::ITEMS_PER_THREAD, + lp::LOAD_ALGORITHM, + lp::LOAD_MODIFIER, + lp::RANK_ALGORITHM, + lp::SCAN_ALGORITHM, + lp::RADIX_BITS}, + SegmentedSortSubWarpMergeSortPolicy{ + mp::BLOCK_THREADS, + mp::WARP_THREADS, + mp::ITEMS_PER_THREAD, + mp::LOAD_ALGORITHM, + mp::LOAD_MODIFIER, + mp::STORE_ALGORITHM}, + SegmentedSortSubWarpMergeSortPolicy{ + sp::BLOCK_THREADS, + sp::WARP_THREADS, + sp::ITEMS_PER_THREAD, + sp::LOAD_ALGORITHM, + sp::LOAD_MODIFIER, + sp::STORE_ALGORITHM}, + ap::PARTITIONING_THRESHOLD}; + } +}; + +// Partition selects large and small groups. The middle group is not selected. +static constexpr size_t num_selected_groups = 2; +} // namespace detail::segmented_sort + +// TODO(bgruber): remove in CCCL 4.0 +//! Deprecated [Since 3.5] +template < + SortOrder Order, + typename KeyT, + typename ValueT, + typename OffsetT, + typename BeginOffsetIteratorT, + typename EndOffsetIteratorT, + typename PolicyHub = detail::segmented_sort::policy_hub, + typename KernelSource = detail::segmented_sort::DeviceSegmentedSortKernelSource< + detail::segmented_sort::policy_selector_from_hub, + Order, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + OffsetT>, + typename PartitionPolicyHub = detail::three_way_partition::policy_hub< + cub::detail::it_value_t>, + detail::three_way_partition::per_partition_offset_t>, + typename PartitionKernelSource = detail::three_way_partition::DeviceThreeWayPartitionKernelSource< + detail::three_way_partition::policy_selector_from_hub, + THRUST_NS_QUALIFIER::counting_iterator, + cub::detail::segmented_sort::local_segment_index_t*, + cub::detail::segmented_sort::local_segment_index_t*, + ::cuda::std::reverse_iterator, + cub::detail::segmented_sort::local_segment_index_t*, + detail::three_way_partition::ScanTileStateT, + cub::detail::segmented_sort::LargeSegmentsSelectorT, + cub::detail::segmented_sort::SmallSegmentsSelectorT, + detail::three_way_partition::per_partition_offset_t, + detail::three_way_partition::streaming_context_t, + detail::choose_signed_offset::type>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceSegmentedSort") DispatchSegmentedSort +{ + using local_segment_index_t = detail::segmented_sort::local_segment_index_t; + using global_segment_offset_t = detail::segmented_sort::global_segment_offset_t; + + static constexpr int KEYS_ONLY = ::cuda::std::is_same_v; + + // Partition selects large and small groups. The middle group is not selected. + static constexpr size_t num_selected_groups = detail::segmented_sort::num_selected_groups; + + /** + * Device-accessible allocation of temporary storage. When `nullptr`, the + * required allocation size is written to `temp_storage_bytes` and no work + * is done. + */ + void* d_temp_storage; + + /// Reference to size in bytes of `d_temp_storage` allocation + size_t& temp_storage_bytes; + + /** + * Double-buffer whose current buffer contains the unsorted input keys and, + * upon return, is updated to point to the sorted output keys + */ + DoubleBuffer& d_keys; + + /** + * Double-buffer whose current buffer contains the unsorted input values and, + * upon return, is updated to point to the sorted output values + */ + DoubleBuffer& d_values; + + /// Number of items to sort + ::cuda::std::int64_t num_items; + + /// The number of segments that comprise the sorting data + global_segment_offset_t num_segments; + + /** + * Random-access input iterator to the sequence of beginning offsets of length + * `num_segments`, such that `d_begin_offsets[i]` is the first element of the + * ith data segment in `d_keys_*` and `d_values_*` + */ + BeginOffsetIteratorT d_begin_offsets; + + /** + * Random-access input iterator to the sequence of ending offsets of length + * `num_segments`, such that d_end_offsets[i]-1 is the last element + * of the ith data segment in `d_keys_*` and + * `d_values_*`. If `d_end_offsets[i]-1 <= d_begin_offsets[i]`, + * the ith is considered empty. + */ + EndOffsetIteratorT d_end_offsets; + + /// Whether is okay to overwrite source buffers + bool is_overwrite_okay; + + /// CUDA stream to launch kernels within. + cudaStream_t stream; + + KernelSource kernel_source; + + PartitionKernelSource partition_kernel_source; + + KernelLauncherFactory launcher_factory; + + typename PartitionPolicyHub::MaxPolicy partition_max_policy; + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT policy = {}) + { + auto wrapped_policy = detail::segmented_sort::MakeSegmentedSortPolicyWrapper(policy); + + CUB_DETAIL_STATIC_ISH_ASSERT(wrapped_policy.LargeSegmentLoadModifier() != CacheLoadModifier::LOAD_LDG, + "The memory consistency model does not apply to texture accesses"); + + CUB_DETAIL_STATIC_ISH_ASSERT( + KEYS_ONLY || wrapped_policy.LargeSegmentLoadAlgorithm() != BLOCK_LOAD_STRIPED + || wrapped_policy.MediumSegmentLoadAlgorithm() != WARP_LOAD_STRIPED + || wrapped_policy.SmallSegmentLoadAlgorithm() != WARP_LOAD_STRIPED, + "Striped load will make this algorithm unstable"); + + CUB_DETAIL_STATIC_ISH_ASSERT(wrapped_policy.MediumSegmentStoreAlgorithm() != WARP_STORE_STRIPED + || wrapped_policy.SmallSegmentStoreAlgorithm() != WARP_STORE_STRIPED, + "Striped stores will produce unsorted results"); + + const int radix_bits = wrapped_policy.LargeSegmentRadixBits(); + + //------------------------------------------------------------------------ + // Prepare temporary storage layout + //------------------------------------------------------------------------ + + const bool partition_segments = num_segments > wrapped_policy.PartitioningThreshold(); + + cub::detail::temporary_storage::layout<5> temporary_storage_layout; + + auto keys_slot = temporary_storage_layout.get_slot(0); + auto values_slot = temporary_storage_layout.get_slot(1); + auto large_and_medium_partitioning_slot = temporary_storage_layout.get_slot(2); + auto small_partitioning_slot = temporary_storage_layout.get_slot(3); + auto group_sizes_slot = temporary_storage_layout.get_slot(4); + + auto keys_allocation = keys_slot->create_alias(); + auto values_allocation = values_slot->create_alias(); + + if (!is_overwrite_okay) + { + keys_allocation.grow(num_items); + + if (!KEYS_ONLY) + { + values_allocation.grow(num_items); + } + } + + auto large_and_medium_segments_indices = large_and_medium_partitioning_slot->create_alias(); + auto small_segments_indices = small_partitioning_slot->create_alias(); + auto group_sizes = group_sizes_slot->create_alias(); + + size_t three_way_partition_temp_storage_bytes{}; + + auto large_segments_selector = + kernel_source.LargeSegmentsSelector(wrapped_policy.MediumPolicyItemsPerTile(), d_begin_offsets, d_end_offsets); + auto small_segments_selector = + kernel_source.SmallSegmentsSelector(wrapped_policy.SmallPolicyItemsPerTile() + 1, d_begin_offsets, d_end_offsets); + + auto device_partition_temp_storage = keys_slot->create_alias(); + + if (partition_segments) + { + constexpr auto num_segments_per_invocation_limit = + static_cast(::cuda::std::numeric_limits::max()); + auto const max_num_segments_per_invocation = static_cast( + (::cuda::std::min) (static_cast(num_segments), num_segments_per_invocation_limit)); + + large_and_medium_segments_indices.grow(max_num_segments_per_invocation); + small_segments_indices.grow(max_num_segments_per_invocation); + group_sizes.grow(num_selected_groups); + + auto medium_indices_iterator = ::cuda::std::make_reverse_iterator(large_and_medium_segments_indices.get()); + + // We call partition through dispatch instead of device because c.parallel needs to be able to call the kernel. + // This approach propagates the type erasure to partition. + using ChooseOffsetT = detail::choose_signed_offset; + using PartitionOffsetT = typename ChooseOffsetT::type; + using DispatchThreeWayPartitionIfT = cub::detail::three_way_partition::dispatch_three_way_partition_if< + THRUST_NS_QUALIFIER::counting_iterator, + decltype(large_and_medium_segments_indices.get()), + decltype(small_segments_indices.get()), + decltype(medium_indices_iterator), + decltype(group_sizes.get()), + decltype(large_segments_selector), + decltype(small_segments_selector), + PartitionOffsetT, + PartitionPolicyHub, + PartitionKernelSource, + KernelLauncherFactory>; + + // Signed integer type for global offsets + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const auto error = ChooseOffsetT::is_exceeding_offset_type(num_items)) + { + return error; + } + + DispatchThreeWayPartitionIfT::Dispatch( + nullptr, + three_way_partition_temp_storage_bytes, + THRUST_NS_QUALIFIER::counting_iterator(0), + large_and_medium_segments_indices.get(), + small_segments_indices.get(), + medium_indices_iterator, + group_sizes.get(), + large_segments_selector, + small_segments_selector, + static_cast(max_num_segments_per_invocation), + stream, + partition_kernel_source, + launcher_factory, + partition_max_policy); + + device_partition_temp_storage.grow(three_way_partition_temp_storage_bytes); + } + + if (d_temp_storage == nullptr) + { + temp_storage_bytes = temporary_storage_layout.get_size(); + + // Return if the caller is simply requesting the size of the storage allocation + return cudaSuccess; + } + + if (num_items == 0 || num_segments == 0) + { + return cudaSuccess; + } + + if (const auto error = CubDebug(temporary_storage_layout.map_to_buffer(d_temp_storage, temp_storage_bytes))) + { + return error; + } + + //------------------------------------------------------------------------ + // Sort + //------------------------------------------------------------------------ + + const bool is_num_passes_odd = GetNumPasses(radix_bits) & 1; + + /** + * This algorithm sorts segments that don't fit into shared memory with + * the in-global-memory radix sort. Radix sort splits key representation + * into multiple "digits". Each digit is RADIX_BITS wide. The algorithm + * iterates over these digits. Each of these iterations consists of a + * couple of stages. The first stage computes a histogram for a current + * digit in each segment key. This histogram helps to determine the + * starting position of the keys group with a similar digit. + * For example: + * keys_digits = [ 1, 0, 0, 1 ] + * digit_prefix = [ 0, 2 ] + * The second stage checks the keys again and increments the prefix to + * determine the final position of the key: + * + * expression | key | idx | result + * ----------------------------------- | ----- | ------- | -------------- + * result[prefix[keys[0]]++] = keys[0] | 1 | 2 | [ ?, ?, 1, ? ] + * result[prefix[keys[1]]++] = keys[0] | 0 | 0 | [ 0, ?, 1, ? ] + * result[prefix[keys[2]]++] = keys[0] | 0 | 1 | [ 0, 0, 1, ? ] + * result[prefix[keys[3]]++] = keys[0] | 1 | 3 | [ 0, 0, 1, 1 ] + * + * If the resulting memory is aliased to the input one, we'll face the + * following issues: + * + * input | key | idx | result/input | issue + * -------------- | ----- | ------- | ---------------- | ---------------- + * [ 1, 0, 0, 1 ] | 1 | 2 | [ 1, 0, 1, 1 ] | overwrite keys[2] + * [ 1, 0, 1, 1 ] | 0 | 0 | [ 0, 0, 1, 1 ] | + * [ 0, 0, 1, 1 ] | 1 | 3 | [ 0, 0, 1, 1 ] | extra key + * [ 0, 0, 1, 1 ] | 1 | 4 | [ 0, 0, 1, 1 ] 1 | OOB access + * + * To avoid these issues, we have to use extra memory. The extra memory + * holds temporary storage for writing intermediate results of each stage. + * Since we iterate over digits in keys, we potentially need: + * `sizeof(KeyT) * num_items * cuda::ceil_div(sizeof(KeyT),RADIX_BITS)` + * auxiliary memory bytes. To reduce the auxiliary memory storage + * requirements, the algorithm relies on a double buffer facility. The + * idea behind it is in swapping destination and source buffers at each + * iteration. This way, we can use only two buffers. One of these buffers + * can be the final algorithm output destination. Therefore, only one + * auxiliary array is needed. Depending on the number of iterations, we + * can initialize the double buffer so that the algorithm output array + * will match the double buffer result one at the final iteration. + * A user can provide this algorithm with a double buffer straightaway to + * further reduce the auxiliary memory requirements. `is_overwrite_okay` + * indicates this use case. + */ + detail::device_double_buffer d_keys_double_buffer( + (is_overwrite_okay || is_num_passes_odd) ? d_keys.Alternate() : keys_allocation.get(), + (is_overwrite_okay) ? d_keys.Current() + : (is_num_passes_odd) ? keys_allocation.get() + : d_keys.Alternate()); + + detail::device_double_buffer d_values_double_buffer( + (is_overwrite_okay || is_num_passes_odd) ? d_values.Alternate() : values_allocation.get(), + (is_overwrite_okay) ? d_values.Current() + : (is_num_passes_odd) ? values_allocation.get() + : d_values.Alternate()); + + cudaError_t error; + if (partition_segments) + { + // Partition input segments into size groups and assign specialized + // kernels for each of them. + error = SortWithPartitioning( + kernel_source.SegmentedSortKernelLarge(), + kernel_source.SegmentedSortKernelSmall(), + three_way_partition_temp_storage_bytes, + d_keys_double_buffer, + d_values_double_buffer, + large_segments_selector, + small_segments_selector, + device_partition_temp_storage, + large_and_medium_segments_indices, + small_segments_indices, + group_sizes, + wrapped_policy); + } + else + { + // If there are not enough segments, there's no reason to spend time + // on extra partitioning steps. + + error = SortWithoutPartitioning( + kernel_source.SegmentedSortFallbackKernel(), d_keys_double_buffer, d_values_double_buffer, wrapped_policy); + } + + d_keys.selector = GetFinalSelector(d_keys.selector, radix_bits); + d_values.selector = GetFinalSelector(d_values.selector, radix_bits); + + return error; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + global_segment_offset_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + bool is_overwrite_okay, + cudaStream_t stream, + KernelSource kernel_source = {}, + PartitionKernelSource partition_kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}, + PartitionMaxPolicyT partition_max_policy = {}) + { + // Get PTX version + int ptx_version = 0; + if (const auto error = CubDebug(launcher_factory.PtxVersion(ptx_version))) + { + return error; + } + + // Create dispatch functor + DispatchSegmentedSort dispatch{ + d_temp_storage, + temp_storage_bytes, + d_keys, + d_values, + num_items, + num_segments, + d_begin_offsets, + d_end_offsets, + is_overwrite_okay, + stream, + kernel_source, + partition_kernel_source, + launcher_factory, + partition_max_policy}; + + // Dispatch to chained policy + return CubDebug(max_policy.Invoke(ptx_version, dispatch)); + } + +private: + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE int GetNumPasses(int radix_bits) + { + constexpr int byte_size = 8; + const int num_bits = static_cast(kernel_source.KeySize()) * byte_size; + const int num_passes = ::cuda::ceil_div(num_bits, radix_bits); + return num_passes; + } + + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE int GetFinalSelector(int selector, int radix_bits) + { + // Sorted data always ends up in the other vector + if (!is_overwrite_okay) + { + return (selector + 1) & 1; + } + + return (selector + GetNumPasses(radix_bits)) & 1; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE T* GetFinalOutput(int radix_bits, DoubleBuffer& buffer) + { + const int final_selector = GetFinalSelector(buffer.selector, radix_bits); + return buffer.d_buffers[final_selector]; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t SortWithPartitioning( + LargeKernelT large_kernel, + SmallKernelT small_kernel, + size_t three_way_partition_temp_storage_bytes, + cub::detail::device_double_buffer& d_keys_double_buffer, + cub::detail::device_double_buffer& d_values_double_buffer, + typename KernelSource::LargeSegmentsSelectorT& large_segments_selector, + typename KernelSource::SmallSegmentsSelectorT& small_segments_selector, + cub::detail::temporary_storage::alias& device_partition_temp_storage, + cub::detail::temporary_storage::alias& large_and_medium_segments_indices, + cub::detail::temporary_storage::alias& small_segments_indices, + cub::detail::temporary_storage::alias& group_sizes, + WrappedPolicyT wrapped_policy) + { + constexpr global_segment_offset_t num_segments_per_invocation_limit = + static_cast(::cuda::std::numeric_limits::max()); + + // We repeatedly invoke the partitioning and sorting kernels until all segments are processed. + const global_segment_offset_t num_invocations = + ::cuda::ceil_div(static_cast(num_segments), num_segments_per_invocation_limit); + for (global_segment_offset_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const global_segment_offset_t current_seg_offset = invocation_index * num_segments_per_invocation_limit; + const local_segment_index_t current_num_segments = + (invocation_index == (num_invocations - 1)) + ? static_cast(num_segments - current_seg_offset) + : num_segments_per_invocation_limit; + + kernel_source.SetSegmentOffset(large_segments_selector, current_seg_offset); + kernel_source.SetSegmentOffset(small_segments_selector, current_seg_offset); + + BeginOffsetIteratorT current_begin_offset = d_begin_offsets; + EndOffsetIteratorT current_end_offset = d_end_offsets; + + current_begin_offset += current_seg_offset; + current_end_offset += current_seg_offset; + + auto medium_indices_iterator = + ::cuda::std::make_reverse_iterator(large_and_medium_segments_indices.get() + current_num_segments); + + // We call partition through dispatch instead of device because c.parallel needs to be able to call the kernel. + // This approach propagates the type erasure to partition. + using ChooseOffsetT = detail::choose_signed_offset; + using PartitionOffsetT = typename ChooseOffsetT::type; + using DispatchThreeWayPartitionIfT = cub::detail::three_way_partition::dispatch_three_way_partition_if< + THRUST_NS_QUALIFIER::counting_iterator, + decltype(large_and_medium_segments_indices.get()), + decltype(small_segments_indices.get()), + decltype(medium_indices_iterator), + decltype(group_sizes.get()), + decltype(large_segments_selector), + decltype(small_segments_selector), + PartitionOffsetT, + PartitionPolicyHub, + PartitionKernelSource, + KernelLauncherFactory>; + + // Signed integer type for global offsets + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const auto error = ChooseOffsetT::is_exceeding_offset_type(num_items)) + { + return error; + } + + if (const auto error = DispatchThreeWayPartitionIfT::Dispatch( + device_partition_temp_storage.get(), + three_way_partition_temp_storage_bytes, + THRUST_NS_QUALIFIER::counting_iterator(0), + large_and_medium_segments_indices.get(), + small_segments_indices.get(), + medium_indices_iterator, + group_sizes.get(), + large_segments_selector, + small_segments_selector, + static_cast(current_num_segments), + stream, + partition_kernel_source, + launcher_factory, + partition_max_policy)) + { + return error; + } + + // The device path is only used (and only compiles) when CDP is enabled. + // It's defined in a macro since we can't put `#ifdef`s inside of + // `NV_IF_TARGET`. +#ifndef CUB_RDC_ENABLED +# define CUB_TEMP_DEVICE_CODE +#else // CUB_RDC_ENABLED +# define CUB_TEMP_DEVICE_CODE \ + if (const auto error = CubDebug( \ + launcher_factory(1, 1, 0, stream) \ + .doit( \ + detail::segmented_sort::DeviceSegmentedSortContinuationKernel< \ + LargeKernelT, \ + SmallKernelT, \ + KeyT, \ + ValueT, \ + BeginOffsetIteratorT, \ + EndOffsetIteratorT, \ + KernelLauncherFactory>, \ + large_kernel, \ + small_kernel, \ + current_num_segments, \ + d_keys.Current(), \ + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_keys), \ + d_keys_double_buffer, \ + d_values.Current(), \ + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_values), \ + d_values_double_buffer, \ + current_begin_offset, \ + current_end_offset, \ + group_sizes.get(), \ + large_and_medium_segments_indices.get(), \ + small_segments_indices.get(), \ + launcher_factory, \ + wrapped_policy.large_segment.ThreadsPerBlock(), \ + wrapped_policy.small_segment.ThreadsPerBlock(), \ + wrapped_policy.medium_segment.SegmentsPerBlock(), \ + wrapped_policy.small_segment.SegmentsPerBlock()))) \ + { \ + return error; \ + } \ + \ + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) \ + { \ + return error; \ + } +#endif // CUB_RDC_ENABLED + + NV_IF_ELSE_TARGET( + NV_IS_HOST, + ({ + local_segment_index_t h_group_sizes[num_selected_groups]; + if (const auto error = CubDebug(launcher_factory.MemcpyAsync( + h_group_sizes, + group_sizes.get(), + num_selected_groups * sizeof(local_segment_index_t), + cudaMemcpyDeviceToHost, + stream))) + { + return error; + } + + if (const auto error = CubDebug(SyncStream(stream))) + { + return error; + } + + if (const auto error = detail::segmented_sort::device_segmented_sort_continuation( + large_kernel, + small_kernel, + current_num_segments, + d_keys.Current(), + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_keys), + d_keys_double_buffer, + d_values.Current(), + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_values), + d_values_double_buffer, + current_begin_offset, + current_end_offset, + h_group_sizes, + large_and_medium_segments_indices.get(), + small_segments_indices.get(), + stream, + launcher_factory, + wrapped_policy.LargeSegmentThreadsPerBlock(), + wrapped_policy.SmallSegmentThreadsPerBlock(), + wrapped_policy.SegmentsPerMediumBlock(), + wrapped_policy.SegmentsPerSmallBlock())) + { + return error; + } + }), + // NV_IS_DEVICE: + (CUB_TEMP_DEVICE_CODE)); + } +#undef CUB_TEMP_DEVICE_CODE + + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t SortWithoutPartitioning( + FallbackKernelT fallback_kernel, + cub::detail::device_double_buffer& d_keys_double_buffer, + cub::detail::device_double_buffer& d_values_double_buffer, + WrappedPolicyT wrapped_policy) + { + const auto blocks_in_grid = static_cast(num_segments); + const auto threads_in_block = static_cast(wrapped_policy.LargeSegmentThreadsPerBlock()); + +// Log kernel configuration +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceSegmentedSortFallbackKernel<<<%d, %d, " + "0, %lld>>>(), %d items per thread, bit_grain %d\n", + blocks_in_grid, + threads_in_block, + (long long) stream, + wrapped_policy.LargeSegmentItemsPerThread(), + wrapped_policy.LargeSegmentRadixBits()); +#endif // CUB_DEBUG_LOG + + // Invoke fallback kernel + launcher_factory(blocks_in_grid, threads_in_block, 0, stream) + .doit(fallback_kernel, + d_keys.Current(), + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_keys), + d_keys_double_buffer, + d_values.Current(), + GetFinalOutput(wrapped_policy.LargeSegmentRadixBits(), d_values), + d_values_double_buffer, + d_begin_offsets, + d_end_offsets); + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + return CubDebug(detail::DebugSyncStream(stream)); + } +}; + +namespace detail::segmented_sort +{ +template +CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t sort_with_partitioning( + LargeKernelT large_kernel, + SmallKernelT small_kernel, + global_segment_offset_t num_segments, + ::cuda::std::int64_t num_items, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream, + size_t three_way_partition_temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + device_double_buffer& d_keys_double_buffer, + device_double_buffer& d_values_double_buffer, + typename KernelSource::LargeSegmentsSelectorT& large_segments_selector, + typename KernelSource::SmallSegmentsSelectorT& small_segments_selector, + temporary_storage::alias& device_partition_temp_storage, + temporary_storage::alias& large_and_medium_segments_indices, + temporary_storage::alias& small_segments_indices, + temporary_storage::alias& group_sizes, + KernelSource& kernel_source, + PartitionKernelSource& partition_kernel_source, + KernelLauncherFactory& launcher_factory, + PartitionPolicySelector partition_policy_selector, + const SegmentedSortPolicy& active_policy, + GetFinalOutputOp&& get_final_output) +{ + constexpr auto num_segments_per_invocation_limit = + static_cast(::cuda::std::numeric_limits::max()); + + const global_segment_offset_t num_invocations = + ::cuda::ceil_div(static_cast(num_segments), num_segments_per_invocation_limit); + for (global_segment_offset_t invocation_index = 0; invocation_index < num_invocations; invocation_index++) + { + const global_segment_offset_t current_seg_offset = invocation_index * num_segments_per_invocation_limit; + const local_segment_index_t current_num_segments = + (invocation_index == (num_invocations - 1)) + ? static_cast(num_segments - current_seg_offset) + : num_segments_per_invocation_limit; + + kernel_source.SetSegmentOffset(large_segments_selector, current_seg_offset); + kernel_source.SetSegmentOffset(small_segments_selector, current_seg_offset); + + BeginOffsetIteratorT current_begin_offset = d_begin_offsets; + EndOffsetIteratorT current_end_offset = d_end_offsets; + current_begin_offset += current_seg_offset; + current_end_offset += current_seg_offset; + + auto medium_indices_iterator = + ::cuda::std::make_reverse_iterator(large_and_medium_segments_indices.get() + current_num_segments); + + using ChooseOffsetT = choose_signed_offset; + if (const auto error = ChooseOffsetT::is_exceeding_offset_type(num_items)) + { + return error; + } + + if (const auto error = three_way_partition::dispatch( + device_partition_temp_storage.get(), + three_way_partition_temp_storage_bytes, + THRUST_NS_QUALIFIER::counting_iterator(0), + large_and_medium_segments_indices.get(), + small_segments_indices.get(), + medium_indices_iterator, + group_sizes.get(), + large_segments_selector, + small_segments_selector, + static_cast(current_num_segments), + stream, + partition_policy_selector, + partition_kernel_source, + launcher_factory)) + { + return error; + } + + [[maybe_unused]] auto device_path = [&] { +#ifdef CUB_RDC_ENABLED + if (const auto error = CubDebug( + launcher_factory(1, 1, 0, stream) + .doit( + detail::segmented_sort::DeviceSegmentedSortContinuationKernel< + LargeKernelT, + SmallKernelT, + KeyT, + ValueT, + BeginOffsetIteratorT, + EndOffsetIteratorT, + KernelLauncherFactory>, + large_kernel, + small_kernel, + current_num_segments, + d_keys.Current(), + get_final_output(d_keys, active_policy.large_segment.radix_bits), + d_keys_double_buffer, + d_values.Current(), + get_final_output(d_values, active_policy.large_segment.radix_bits), + d_values_double_buffer, + current_begin_offset, + current_end_offset, + group_sizes.get(), + large_and_medium_segments_indices.get(), + small_segments_indices.get(), + launcher_factory, + active_policy.large_segment.threads_per_block, + active_policy.small_segment.threads_per_block, + active_policy.medium_segment.segments_per_block(), + active_policy.small_segment.segments_per_block()))) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } +#endif // CUB_RDC_ENABLED + return cudaSuccess; + }; + + NV_IF_ELSE_TARGET( + NV_IS_HOST, + ({ + local_segment_index_t h_group_sizes[num_selected_groups]; + if (const auto error = CubDebug(launcher_factory.MemcpyAsync( + h_group_sizes, + group_sizes.get(), + num_selected_groups * sizeof(local_segment_index_t), + cudaMemcpyDeviceToHost, + stream))) + { + return error; + } + + if (const auto error = CubDebug(SyncStream(stream))) + { + return error; + } + + if (const auto error = detail::segmented_sort::device_segmented_sort_continuation( + large_kernel, + small_kernel, + current_num_segments, + d_keys.Current(), + get_final_output(d_keys, active_policy.large_segment.radix_bits), + d_keys_double_buffer, + d_values.Current(), + get_final_output(d_values, active_policy.large_segment.radix_bits), + d_values_double_buffer, + current_begin_offset, + current_end_offset, + h_group_sizes, + large_and_medium_segments_indices.get(), + small_segments_indices.get(), + stream, + launcher_factory, + active_policy.large_segment.threads_per_block, + active_policy.small_segment.threads_per_block, + active_policy.medium_segment.segments_per_block(), + active_policy.small_segment.segments_per_block())) + { + return error; + } + }), + ({ + if (const auto error = device_path()) + { + return error; + } + })); + } + + return cudaSuccess; +} + +template +CUB_RUNTIME_FUNCTION void check_policy(PolicyGetter policy_getter) +{ + [[maybe_unused]] CUB_DETAIL_CONSTEXPR_ISH SegmentedSortPolicy active_policy = policy_getter(); + + CUB_DETAIL_STATIC_ISH_ASSERT(active_policy.large_segment.load_modifier != CacheLoadModifier::LOAD_LDG, + "The memory consistency model does not apply to texture accesses"); + CUB_DETAIL_STATIC_ISH_ASSERT( + KeysOnly || active_policy.large_segment.load_algorithm != BLOCK_LOAD_STRIPED + || active_policy.medium_segment.load_algorithm != WARP_LOAD_STRIPED + || active_policy.small_segment.load_algorithm != WARP_LOAD_STRIPED, + "Striped load will make this algorithm unstable"); + CUB_DETAIL_STATIC_ISH_ASSERT(active_policy.medium_segment.store_algorithm != WARP_STORE_STRIPED + || active_policy.small_segment.store_algorithm != WARP_STORE_STRIPED, + "Striped stores will produce unsorted results"); +} + +template +CUB_RUNTIME_FUNCTION _CCCL_VISIBILITY_HIDDEN _CCCL_FORCEINLINE cudaError_t sort_without_partitioning( + FallbackKernelT fallback_kernel, + global_segment_offset_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + cudaStream_t stream, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + device_double_buffer& d_keys_double_buffer, + device_double_buffer& d_values_double_buffer, + KernelLauncherFactory& launcher_factory, + const SegmentedSortPolicy& active_policy, + GetFinalOutputOp&& get_final_output) +{ + const auto blocks_in_grid = static_cast(num_segments); + const auto threads_in_block = static_cast(active_policy.large_segment.threads_per_block); +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking DeviceSegmentedSortFallbackKernel<<<%d, %d, 0, %lld>>>(), %d items per thread, bit_grain %d\n", + blocks_in_grid, + threads_in_block, + (long long) stream, + active_policy.large_segment.items_per_thread, + active_policy.large_segment.radix_bits); +#endif // CUB_DEBUG_LOG + + if (const auto error = CubDebug( + launcher_factory(blocks_in_grid, threads_in_block, 0, stream) + .doit(fallback_kernel, + d_keys.Current(), + get_final_output(d_keys, active_policy.large_segment.radix_bits), + d_keys_double_buffer, + d_values.Current(), + get_final_output(d_values, active_policy.large_segment.radix_bits), + d_values_double_buffer, + d_begin_offsets, + d_end_offsets))) + { + return error; + } + + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + return cudaSuccess; +} + +template < + SortOrder Order, + typename OffsetT, + typename KeyT, + typename ValueT, + typename BeginOffsetIteratorT, + typename EndOffsetIteratorT, + typename PolicySelector = policy_selector_from_types, + typename KernelSource = + DeviceSegmentedSortKernelSource, + typename PartitionPolicySelector = detail::three_way_partition::policy_selector_from_types< + cub::detail::it_value_t>, + three_way_partition::per_partition_offset_t>, + typename PartitionKernelSource = detail::three_way_partition::DeviceThreeWayPartitionKernelSource< + PartitionPolicySelector, + THRUST_NS_QUALIFIER::counting_iterator, + local_segment_index_t*, + local_segment_index_t*, + ::cuda::std::reverse_iterator, + local_segment_index_t*, + three_way_partition::ScanTileStateT, + LargeSegmentsSelectorT, + SmallSegmentsSelectorT, + three_way_partition::per_partition_offset_t, + three_way_partition::streaming_context_t, + choose_signed_offset::type>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires segmented_sort_policy_selector + && three_way_partition::three_way_partition_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + DoubleBuffer& d_keys, + DoubleBuffer& d_values, + ::cuda::std::int64_t num_items, + global_segment_offset_t num_segments, + BeginOffsetIteratorT d_begin_offsets, + EndOffsetIteratorT d_end_offsets, + bool is_overwrite_okay, + cudaStream_t stream, + PolicySelector policy_selector = {}, + PartitionPolicySelector partition_policy_selector = {}, + KernelSource kernel_source = {}, + PartitionKernelSource partition_kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) -> cudaError_t +{ + [[maybe_unused]] static constexpr bool keys_only = ::cuda::std::is_same_v; + + const auto get_num_passes = [&](int radix_bits) { + const int num_bits = static_cast(kernel_source.KeySize()) * CHAR_BIT; + const int num_passes = ::cuda::ceil_div(num_bits, radix_bits); + return num_passes; + }; + + const auto get_final_selector = [&](int selector, int radix_bits) { + if (!is_overwrite_okay) + { + return (selector + 1) & 1; + } + return (selector + get_num_passes(radix_bits)) & 1; + }; + + const auto get_final_output = [&](auto& buffer, int radix_bits) { + const int final_selector = get_final_selector(buffer.selector, radix_bits); + return buffer.d_buffers[final_selector]; + }; + + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + return detail::dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) -> cudaError_t { + check_policy(policy_getter); // MSVC fails to evaluate static_asserts inside this lambda, so move them to + // a function + CUB_DETAIL_CONSTEXPR_ISH const SegmentedSortPolicy active_policy = policy_getter(); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceSegmentedSort to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + const int radix_bits = active_policy.large_segment.radix_bits; + + //------------------------------------------------------------------------ + // Prepare temporary storage layout + //------------------------------------------------------------------------ + + const bool partition_segments = num_segments > active_policy.partitioning_threshold; + + temporary_storage::layout<5> temporary_storage_layout; + auto keys_slot = temporary_storage_layout.get_slot(0); + auto values_slot = temporary_storage_layout.get_slot(1); + auto large_and_medium_partitioning_slot = temporary_storage_layout.get_slot(2); + auto small_partitioning_slot = temporary_storage_layout.get_slot(3); + auto group_sizes_slot = temporary_storage_layout.get_slot(4); + + auto keys_allocation = keys_slot->create_alias(); + auto values_allocation = values_slot->create_alias(); + + if (!is_overwrite_okay) + { + keys_allocation.grow(num_items); + if (!keys_only) + { + values_allocation.grow(num_items); + } + } + + auto large_and_medium_segments_indices = large_and_medium_partitioning_slot->create_alias(); + auto small_segments_indices = small_partitioning_slot->create_alias(); + auto group_sizes = group_sizes_slot->create_alias(); + + size_t three_way_partition_temp_storage_bytes{}; + + auto large_segments_selector = kernel_source.LargeSegmentsSelector( + active_policy.medium_segment.items_per_tile(), d_begin_offsets, d_end_offsets); + auto small_segments_selector = kernel_source.SmallSegmentsSelector( + static_cast(active_policy.small_segment.items_per_tile()) + 1, d_begin_offsets, d_end_offsets); + + auto device_partition_temp_storage = keys_slot->create_alias(); + if (partition_segments) + { + constexpr auto num_segments_per_invocation_limit = + static_cast(::cuda::std::numeric_limits::max()); + auto const max_num_segments_per_invocation = static_cast( + (::cuda::std::min) (static_cast(num_segments), num_segments_per_invocation_limit)); + + large_and_medium_segments_indices.grow(max_num_segments_per_invocation); + small_segments_indices.grow(max_num_segments_per_invocation); + group_sizes.grow(num_selected_groups); + + auto medium_indices_iterator = ::cuda::std::make_reverse_iterator(large_and_medium_segments_indices.get()); + + // We call partition through dispatch instead of device because c.parallel needs to be able to call the kernel. + // This approach propagates the type erasure to partition. + using ChooseOffsetT = choose_signed_offset; + + // Signed integer type for global offsets + // Check if the number of items exceeds the range covered by the selected signed offset type + if (const auto error = ChooseOffsetT::is_exceeding_offset_type(num_items)) + { + return error; + } + + three_way_partition::dispatch( + nullptr, + three_way_partition_temp_storage_bytes, + THRUST_NS_QUALIFIER::counting_iterator(0), + large_and_medium_segments_indices.get(), + small_segments_indices.get(), + medium_indices_iterator, + group_sizes.get(), + large_segments_selector, + small_segments_selector, + static_cast(max_num_segments_per_invocation), + stream, + partition_policy_selector, + partition_kernel_source, + launcher_factory); + + device_partition_temp_storage.grow(three_way_partition_temp_storage_bytes); + } + + if (d_temp_storage == nullptr) + { + temp_storage_bytes = temporary_storage_layout.get_size(); + + // Return if the caller is simply requesting the size of the storage allocation + return cudaSuccess; + } + + if (num_items == 0 || num_segments == 0) + { + return cudaSuccess; + } + + if (const auto error = CubDebug(temporary_storage_layout.map_to_buffer(d_temp_storage, temp_storage_bytes))) + { + return error; + } + + //------------------------------------------------------------------------ + // Sort + //------------------------------------------------------------------------ + + const bool is_num_passes_odd = get_num_passes(radix_bits) & 1; + + // This algorithm sorts segments that don't fit into shared memory with the in-global-memory radix sort. Radix sort + // splits key representation into multiple "digits". Each digit is RADIX_BITS wide. The algorithm iterates over + // these digits. Each of these iterations consists of a couple of stages. The first stage computes a histogram for a + // current digit in each segment key. This histogram helps to determine the starting position of the keys group with + // a similar digit. + // For example: + // keys_digits = [ 1, 0, 0, 1 ] + // digit_prefix = [ 0, 2 ] + // The second stage checks the keys again and increments the prefix to determine the final position of the key: + // + // expression | key | idx | result + // ----------------------------------- | ----- | ------- | -------------- + // result[prefix[keys[0]]++] = keys[0] | 1 | 2 | [ ?, ?, 1, ? ] + // result[prefix[keys[1]]++] = keys[0] | 0 | 0 | [ 0, ?, 1, ? ] + // result[prefix[keys[2]]++] = keys[0] | 0 | 1 | [ 0, 0, 1, ? ] + // result[prefix[keys[3]]++] = keys[0] | 1 | 3 | [ 0, 0, 1, 1 ] + // + // If the resulting memory is aliased to the input one, we'll face the following issues: + // + // input | key | idx | result/input | issue + // -------------- | ----- | ------- | ---------------- | ---------------- + // [ 1, 0, 0, 1 ] | 1 | 2 | [ 1, 0, 1, 1 ] | overwrite keys[2] + // [ 1, 0, 1, 1 ] | 0 | 0 | [ 0, 0, 1, 1 ] | + // [ 0, 0, 1, 1 ] | 1 | 3 | [ 0, 0, 1, 1 ] | extra key + // [ 0, 0, 1, 1 ] | 1 | 4 | [ 0, 0, 1, 1 ] 1 | OOB access + // + // To avoid these issues, we have to use extra memory. The extra memory holds temporary storage for writing + // intermediate results of each stage. Since we iterate over digits in keys, we potentially need: `sizeof(KeyT) * + // num_items * cuda::ceil_div(sizeof(KeyT),RADIX_BITS)` auxiliary memory bytes. To reduce the auxiliary memory + // storage requirements, the algorithm relies on a double buffer facility. The idea behind it is in swapping + // destination and source buffers at each iteration. This way, we can use only two buffers. One of these buffers can + // be the final algorithm output destination. Therefore, only one auxiliary array is needed. Depending on the number + // of iterations, we can initialize the double buffer so that the algorithm output array will match the double + // buffer result one at the final iteration. A user can provide this algorithm with a double buffer straightaway to + // further reduce the auxiliary memory requirements. `is_overwrite_okay` + // indicates this use case. + + device_double_buffer d_keys_double_buffer( + (is_overwrite_okay || is_num_passes_odd) ? d_keys.Alternate() : keys_allocation.get(), + (is_overwrite_okay) ? d_keys.Current() + : (is_num_passes_odd) ? keys_allocation.get() + : d_keys.Alternate()); + device_double_buffer d_values_double_buffer( + (is_overwrite_okay || is_num_passes_odd) ? d_values.Alternate() : values_allocation.get(), + (is_overwrite_okay) ? d_values.Current() + : (is_num_passes_odd) ? values_allocation.get() + : d_values.Alternate()); + + const auto segmented_sort_fallback_kernel = kernel_source.SegmentedSortFallbackKernel(); + const auto segmented_sort_kernel_small = kernel_source.SegmentedSortKernelSmall(); + const auto segmented_sort_kernel_large = kernel_source.SegmentedSortKernelLarge(); + + if (partition_segments) + { + if (const auto error = sort_with_partitioning( + segmented_sort_kernel_large, + segmented_sort_kernel_small, + num_segments, + num_items, + d_begin_offsets, + d_end_offsets, + stream, + three_way_partition_temp_storage_bytes, + d_keys, + d_values, + d_keys_double_buffer, + d_values_double_buffer, + large_segments_selector, + small_segments_selector, + device_partition_temp_storage, + large_and_medium_segments_indices, + small_segments_indices, + group_sizes, + kernel_source, + partition_kernel_source, + launcher_factory, + partition_policy_selector, + active_policy, + get_final_output)) + { + return error; + } + } + else + { + if (const auto error = sort_without_partitioning( + segmented_sort_fallback_kernel, + num_segments, + d_begin_offsets, + d_end_offsets, + stream, + d_keys, + d_values, + d_keys_double_buffer, + d_values_double_buffer, + launcher_factory, + active_policy, + get_final_output)) + { + return error; + } + } + + d_keys.selector = get_final_selector(d_keys.selector, radix_bits); + d_values.selector = get_final_selector(d_values.selector, radix_bits); + return cudaSuccess; + }); +} +} // namespace detail::segmented_sort + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_select_if.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_select_if.cuh new file mode 100644 index 00000000..fa8bc5c8 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_select_if.cuh @@ -0,0 +1,1145 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved. +// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +/** + * @file + * cub::DeviceSelect provides device-wide, parallel operations for selecting items from sequences + * of data items residing within device-accessible memory. + */ + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +_CCCL_DIAG_PUSH +_CCCL_DIAG_SUPPRESS_GCC("-Wattributes") // __visibility__ attribute ignored +_CCCL_DIAG_SUPPRESS_NVHPC(attribute_requires_external_linkage) + +CUB_NAMESPACE_BEGIN + +namespace detail::select +{ +// Offset type used to instantiate the stream compaction-kernel and agent to index the items within one partition +using per_partition_offset_t = ::cuda::std::int32_t; + +template +class streaming_context_t +{ +private: + bool first_partition = true; + bool last_partition = false; + TotalNumItemsT total_num_items{}; + TotalNumItemsT total_previous_num_items{}; + + // We use a double-buffer for keeping track of the number of previously selected items + TotalNumItemsT* d_num_selected_in = nullptr; + TotalNumItemsT* d_num_selected_out = nullptr; + +public: + using total_num_items_t = TotalNumItemsT; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE streaming_context_t( + TotalNumItemsT* d_num_selected_in, + TotalNumItemsT* d_num_selected_out, + TotalNumItemsT total_num_items, + bool is_last_partition) + : last_partition(is_last_partition) + , total_num_items(total_num_items) + , d_num_selected_in(d_num_selected_in) + , d_num_selected_out(d_num_selected_out) + {} + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(TotalNumItemsT num_items, bool next_partition_is_the_last) + { + using ::cuda::std::swap; + swap(d_num_selected_in, d_num_selected_out); + first_partition = false; + last_partition = next_partition_is_the_last; + total_previous_num_items += num_items; + }; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT input_offset() const + { + return first_partition ? TotalNumItemsT{0} : total_previous_num_items; + }; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT is_first_partition() const + { + return first_partition; + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected() const + { + return first_partition ? TotalNumItemsT{0} : *d_num_selected_in; + }; + + _CCCL_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_rejected() const + { + return first_partition ? TotalNumItemsT{0} : (total_previous_num_items - num_previously_selected()); + }; + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_total_items(OffsetT) const + { + return total_num_items; + } + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void + update_num_selected(NumSelectedIteratorT user_num_selected_out_it, OffsetT num_selections) const + { + if (last_partition) + { + *user_num_selected_out_it = num_previously_selected() + static_cast(num_selections); + } + else + { + *d_num_selected_out = num_previously_selected() + static_cast(num_selections); + } + } +}; + +template +class streaming_context_t +{ +public: + using total_num_items_t = TotalNumItemsT; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE streaming_context_t(TotalNumItemsT*, TotalNumItemsT*, TotalNumItemsT, bool) {} + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(TotalNumItemsT, bool) {}; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT input_offset() const + { + return TotalNumItemsT{0}; + }; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT is_first_partition() const + { + return true; + }; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_selected() const + { + return TotalNumItemsT{0}; + }; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_previously_rejected() const + { + return TotalNumItemsT{0}; + }; + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE TotalNumItemsT num_total_items(OffsetT num_partition_items) const + { + return num_partition_items; + } + + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void + update_num_selected(NumSelectedIteratorT user_num_selected_out_it, OffsetT num_selections) const + { + *user_num_selected_out_it = num_selections; + } +}; + +/** + * @brief Wrapper that partially specializes the `AgentSelectIf` on the non-type name parameter `KeepRejects`. + */ +template +struct bind_selection_opt +{ + // Using an explicit list of template parameters forwarded to AgentSelectIf, since MSVC complains about a template + // argument following a parameter pack expansion like `AgentSelectIf` + template + using agent_t = + AgentSelectIf; +}; + +template +struct make_vsmem_helper +{ + static constexpr SelectPolicy active_policy = DefaultPolicyGetter{}(); + using agent_policy_t = detail::agent_select_if_policy< + active_policy.lookback.threads_per_block, + active_policy.lookback.items_per_thread, + active_policy.lookback.load_algorithm, + active_policy.lookback.load_modifier, + active_policy.lookback.scan_algorithm, + delay_constructor_t>; + using type = vsmem_helper_default_fallback_policy_t< + agent_policy_t, + bind_selection_opt::template agent_t, + InputIteratorT, + FlagsInputIteratorT, + SelectedOutputIteratorT, + SelectOpT, + EqualityOpT, + OffsetT, + StreamingContextT>; +}; + +/****************************************************************************** + * Kernel entry points + *****************************************************************************/ + +/** + * Select kernel entry point (multi-block) + * + * Performs functor-based selection if SelectOpT functor type != NullType + * Otherwise performs flag-based selection if FlagsInputIterator's value type != NullType + * Otherwise performs discontinuity selection (keep unique) + * + * @tparam PolicySelectorT + * Selects the tuning policy + * + * @tparam InputIteratorT + * Random-access input iterator type for reading input items + * + * @tparam FlagsInputIteratorT + * Random-access input iterator type for reading selection flags (NullType* if a selection functor + * or discontinuity flagging is to be used for selection) + * + * @tparam SelectedOutputIteratorT + * Random-access output iterator type for writing selected items + * + * @tparam NumSelectedIteratorT + * Output iterator type for recording the number of items selected + * + * @tparam ScanTileStateT + * Tile status interface type + * + * @tparam SelectOpT + * Selection operator type (NullType if selection flags or discontinuity flagging is + * to be used for selection) + * + * @tparam EqualityOpT + * Equality operator type (NullType if selection functor or selection flags is + * to be used for selection) + * + * @tparam OffsetT + * Signed integer type for offsets within a partition + * + * @tparam StreamingContextT + * Type providing the context information for the current partition, with the following member functions: + * input_offset() -> base offset for the input (and flags) iterator + * num_previously_selected() -> base offset for the output iterator for selected items + * num_previously_rejected() -> base offset for the output iterator for rejected items (partition only) + * num_total_items() -> total number of items across all partitions (partition only) + * update_num_selected(d_num_sel_out, num_selected) -> invoked by last CTA with number of selected + * + * @tparam KeepRejects + * Whether or not we push rejected items to the back of the output + * + * @param[in] d_in + * Pointer to the input sequence of data items + * + * @param[in] d_flags + * Pointer to the input sequence of selection flags (if applicable) + * + * @param[out] d_selected_out + * Pointer to the output sequence of selected data items + * + * @param[out] d_num_selected_out + * Pointer to the total number of items selected (i.e., length of \p d_selected_out) + * + * @param[in] tile_status + * Tile status interface + * + * @param[in] select_op + * Selection operator + * + * @param[in] equality_op + * Equality operator + * + * @param[in] num_items + * Total number of input items (i.e., length of \p d_in) + * + * @param[in] num_tiles + * Total number of tiles for the entire problem + * + * @param[in] streaming_context + * The context information for the current partition + * + * @param[in] vsmem + * Memory to support virtual shared memory + */ +template +#if _CCCL_HAS_CONCEPTS() + requires select_if_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int( + make_vsmem_helper, + SelectionOpt, + InputIteratorT, + FlagsInputIteratorT, + SelectedOutputIteratorT, + SelectOpT, + EqualityOpT, + OffsetT, + StreamingContextT>::type::agent_policy_t::BLOCK_THREADS)) + _CCCL_KERNEL_ATTRIBUTES void DeviceSelectSweepKernel( + const InputIteratorT d_in, + const FlagsInputIteratorT d_flags, + const SelectedOutputIteratorT d_selected_out, + const NumSelectedIteratorT d_num_selected_out, + ScanTileStateT tile_status, + SelectOpT select_op, + EqualityOpT equality_op, + const OffsetT num_items, + const int num_tiles, + _CCCL_GRID_CONSTANT const StreamingContextT streaming_context, + vsmem_t vsmem) +{ + using VsmemHelperT = typename make_vsmem_helper< + device_policy_getter, + SelectionOpt, + InputIteratorT, + FlagsInputIteratorT, + SelectedOutputIteratorT, + SelectOpT, + EqualityOpT, + OffsetT, + StreamingContextT>::type; + + // Thread block type for selecting data from input tiles + using AgentSelectIfT = typename VsmemHelperT::agent_t; + + // Static shared memory allocation + __shared__ typename VsmemHelperT::static_temp_storage_t static_temp_storage; + + // Get temporary storage + typename AgentSelectIfT::TempStorage& temp_storage = VsmemHelperT::get_temp_storage(static_temp_storage, vsmem); + + // Process tiles + AgentSelectIfT(temp_storage, d_in, d_flags, d_selected_out, select_op, equality_op, num_items, streaming_context) + .ConsumeRange(num_tiles, tile_status, d_num_selected_out); + + // If applicable, hints to discard modified cache lines for vsmem + VsmemHelperT::discard_temp_storage(temp_storage); +} + +// TODO(bgruber): remove in CCCL 4.0 +template +struct policy_selector_from_hub +{ + // this is only called in device code + [[nodiscard]] _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability /*cc*/) const -> SelectPolicy + { + using active_policy = typename PolicyHub::MaxPolicy::ActivePolicy::SelectIfPolicyT; + return SelectPolicy{ + SelectAlgorithm::lookback, + {active_policy::BLOCK_THREADS, + active_policy::ITEMS_PER_THREAD, + active_policy::LOAD_ALGORITHM, + active_policy::LOAD_MODIFIER, + active_policy::SCAN_ALGORITHM, + lookback_delay_policy_from_type}}; + } +}; +} // namespace detail::select + +/****************************************************************************** + * Dispatch + ******************************************************************************/ + +/** + * Utility class for dispatching the appropriately-tuned kernels for DeviceSelect and DevicePartition + * + * Deprecated [Since 3.5] + * + * @tparam InputIteratorT + * Random-access input iterator type for reading input items + * + * @tparam FlagsInputIteratorT + * Random-access input iterator type for reading selection flags (NullType* if a selection functor or discontinuity + * flagging is used for selection) + * + * @tparam SelectedOutputIteratorT + * Random-access output iterator type for writing selected items + * + * @tparam NumSelectedIteratorT + * Output iterator type for recording the number of items selected + * + * @tparam SelectOpT + * Selection operator type (NullType if selection flags or discontinuity flagging is used for selection) + * + * @tparam EqualityOpT + * Equality operator type (NullType if selection functor or selection flags are used for selection) + * + * @tparam OffsetT + * Signed integer type for global offsets + * + * @tparam SelectionOpt + * SelectImpl indicating whether to partition, just selection or selection where the memory for the input and + * output may alias each other. + */ +template < + typename InputIteratorT, + typename FlagsInputIteratorT, + typename SelectedOutputIteratorT, + typename NumSelectedIteratorT, + typename SelectOpT, + typename EqualityOpT, + typename OffsetT, + SelectImpl SelectionOpt, + typename PolicyHub = detail::select::policy_hub< + detail::it_value_t, + detail::it_value_t, + // if/flagged/unique only have a single code path for different offset types, partition has different code paths + ::cuda::std::conditional_t, + detail::select::is_partition_distinct_output_t::value, + SelectionOpt>> +struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceSelect/DevicePartition") DispatchSelectIf +{ + /****************************************************************************** + * Types and constants + ******************************************************************************/ + + // Offset type used to instantiate the stream compaction-kernel and agent to index the items within one partition + using per_partition_offset_t = detail::select::per_partition_offset_t; + + // Offset type large enough to represent any index within the input and output iterators + using num_total_items_t = OffsetT; + + // Whether the algorithm is a partitioning invocation (versus a selection invocation) + static constexpr bool is_partitioning_invocation = (SelectionOpt == SelectImpl::Partition); + + // We always use a streaming context for selection. However, for a partitioning invocation, we only use a streaming + // context when necessary. I.e., if the values representable by OffsetT exceed the values representable by + // per_partition_offset_t. + static constexpr bool use_streaming_context = + (!is_partitioning_invocation) + || (static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max()) + < static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())); + + using streaming_context_t = detail::select::streaming_context_t; + + using ScanTileStateT = ScanTileState; + + static constexpr int INIT_KERNEL_THREADS = 128; + + /// Device-accessible allocation of temporary storage. + /// When `nullptr`, the required allocation size is written to `temp_storage_bytes` + /// and no work is done. + void* d_temp_storage; + + /// Reference to size in bytes of `d_temp_storage` allocation + size_t& temp_storage_bytes; + + /// Pointer to the input sequence of data items + InputIteratorT d_in; + + /// Pointer to the input sequence of selection flags (if applicable) + FlagsInputIteratorT d_flags; + + /// Pointer to the output sequence of selected data items + SelectedOutputIteratorT d_selected_out; + + /// Pointer to the total number of items selected (i.e., length of `d_selected_out`) + NumSelectedIteratorT d_num_selected_out; + + /// Selection operator + SelectOpT select_op; + + /// Equality operator + EqualityOpT equality_op; + + /// Total number of input items (i.e., length of `d_in`) + OffsetT num_items; + + /// CUDA stream to launch kernels within. Default is stream0. + cudaStream_t stream; + + int ptx_version; + + /** + * @param 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 temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param d_in + * Pointer to the input sequence of data items + * + * @param d_flags + * Pointer to the input sequence of selection flags (if applicable) + * + * @param d_selected_out + * Pointer to the output sequence of selected data items + * + * @param d_num_selected_out + * Pointer to the total number of items selected (i.e., length of `d_selected_out`) + * + * @param select_op + * Selection operator + * + * @param equality_op + * Equality operator + * + * @param num_items + * Total number of input items (i.e., length of `d_in`) + * + * @param stream + * CUDA stream to launch kernels within. Default is stream0. + */ + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchSelectIf( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FlagsInputIteratorT d_flags, + SelectedOutputIteratorT d_selected_out, + NumSelectedIteratorT d_num_selected_out, + SelectOpT select_op, + EqualityOpT equality_op, + OffsetT num_items, + cudaStream_t stream, + int ptx_version) + : d_temp_storage(d_temp_storage) + , temp_storage_bytes(temp_storage_bytes) + , d_in(d_in) + , d_flags(d_flags) + , d_selected_out(d_selected_out) + , d_num_selected_out(d_num_selected_out) + , select_op(select_op) + , equality_op(equality_op) + , num_items(num_items) + , stream(stream) + , ptx_version(ptx_version) + {} + + /****************************************************************************** + * Dispatch entrypoints + ******************************************************************************/ + + /** + * Internal dispatch routine for computing a device-wide selection using the + * specified kernel functions. + */ + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t + Invoke(ScanInitKernelPtrT scan_init_kernel, SelectIfKernelPtrT select_if_kernel) + { + using Policy = typename ActivePolicyT::SelectIfPolicyT; + + using VsmemHelperT = cub::detail::vsmem_helper_default_fallback_policy_t< + Policy, + detail::select::bind_selection_opt::template agent_t, + InputIteratorT, + FlagsInputIteratorT, + SelectedOutputIteratorT, + SelectOpT, + EqualityOpT, + per_partition_offset_t, + streaming_context_t>; + cudaError error = cudaSuccess; + + constexpr auto threads_per_block = VsmemHelperT::agent_policy_t::BLOCK_THREADS; + constexpr auto items_per_thread = VsmemHelperT::agent_policy_t::ITEMS_PER_THREAD; + constexpr auto tile_size = OffsetT{threads_per_block * items_per_thread}; + + // The maximum number of items per partition + static constexpr auto max_supported_partition_size = ::cuda::std::numeric_limits::max(); + static constexpr auto full_tile_partition_size = + max_supported_partition_size - (max_supported_partition_size % (threads_per_block * items_per_thread)); + + // For partitioning invocations, we cap the partition size to the maximum number of items supported. + // For selection invocations, we cap at the largest multiple of a full tile. There's a selection-specific bug where + // we would otherwise overflow indices for the last partial tile, when discounting for the out-of-bounds items. + static constexpr per_partition_offset_t capped_partition_size = + is_partitioning_invocation ? max_supported_partition_size : full_tile_partition_size; + + // The maximum number of items for which we will ever invoke the kernel (i.e. largest partition size) + // The extra check of use_streaming_context ensures that OffsetT is larger than per_partition_offset_t to avoid + // truncation of partition_size + auto const max_partition_size = + (use_streaming_context && num_items > static_cast(capped_partition_size)) + ? static_cast(capped_partition_size) + : num_items; + + // The number of partitions required to "iterate" over the total input (ternary to avoid div-by-zero) + auto const num_partitions = + (max_partition_size == 0) ? static_cast(1) : ::cuda::ceil_div(num_items, max_partition_size); + + // The maximum number of tiles for which we will ever invoke the kernel + auto const max_num_tiles_per_invocation = static_cast(::cuda::ceil_div(max_partition_size, tile_size)); + + // The amount of virtual shared memory to allocate + const auto vsmem_size = max_num_tiles_per_invocation * VsmemHelperT::vsmem_per_block; + + do + { + // Specify temporary storage allocation requirements + ::cuda::std::size_t streaming_selection_storage_bytes = + (num_partitions > 1) ? 2 * sizeof(num_total_items_t) : ::cuda::std::size_t{0}; + ::cuda::std::size_t allocation_sizes[3] = {0ULL, vsmem_size, streaming_selection_storage_bytes}; + + // Bytes needed for tile status descriptors + error = + CubDebug(ScanTileStateT::AllocationSize(static_cast(max_num_tiles_per_invocation), allocation_sizes[0])); + if (cudaSuccess != error) + { + break; + } + + // Compute allocation pointers into the single storage blob (or compute the necessary size of the blob) + void* allocations[3] = {}; + + error = CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)); + if (cudaSuccess != error) + { + break; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage allocation + break; + } + + // Initialize the streaming context with the temporary storage for double-buffering the previously selected items + // and the total number (across all partitions) of items + num_total_items_t* tmp_num_selected_out = reinterpret_cast(allocations[2]); + streaming_context_t streaming_context{ + tmp_num_selected_out, (tmp_num_selected_out + 1), num_items, (num_partitions <= 1)}; + + // Iterate over the partitions until all input is processed + for (OffsetT partition_idx = 0; partition_idx < num_partitions; partition_idx++) + { + OffsetT current_partition_offset = partition_idx * max_partition_size; + OffsetT current_num_items = + (partition_idx + 1 == num_partitions) ? (num_items - current_partition_offset) : max_partition_size; + + // Construct the tile status interface + const auto current_num_tiles = static_cast(::cuda::ceil_div(current_num_items, tile_size)); + ScanTileStateT tile_status; + error = CubDebug(tile_status.Init(current_num_tiles, allocations[0], allocation_sizes[0])); + if (cudaSuccess != error) + { + return error; + } + + // Log scan_init_kernel configuration + int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(current_num_tiles, INIT_KERNEL_THREADS)); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking scan_init_kernel<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + INIT_KERNEL_THREADS, + (long long) stream); +#endif + + // Invoke scan_init_kernel to initialize tile descriptors + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, INIT_KERNEL_THREADS, 0, stream) + .doit(scan_init_kernel, tile_status, current_num_tiles, d_num_selected_out)); + if (cudaSuccess != error) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(detail::DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + + // No more items to process (note, we do not want to return early for num_items==0, because we need to make sure + // that `scan_init_kernel` has written '0' to d_num_selected_out) + if (current_num_items == 0) + { + return cudaSuccess; + } + +// Log select_if_kernel configuration +#ifdef CUB_DEBUG_LOG + { + // Get SM occupancy for select_if_kernel + int range_select_sm_occupancy; + error = CubDebug(MaxSmOccupancy(range_select_sm_occupancy, // out + select_if_kernel, + threads_per_block)); + if (cudaSuccess != error) + { + return error; + } + + _CubLog("Invoking select_if_kernel<<<%d, %d, 0, " + "%lld>>>(), %d items per thread, %d SM occupancy\n", + current_num_tiles, + threads_per_block, + (long long) stream, + items_per_thread, + range_select_sm_occupancy); + } +#endif + + // Invoke select_if_kernel + error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(current_num_tiles, threads_per_block, 0, stream) + .doit(select_if_kernel, + d_in, + d_flags, + d_selected_out, + d_num_selected_out, + tile_status, + select_op, + equality_op, + static_cast(current_num_items), + current_num_tiles, + streaming_context, + cub::detail::vsmem_t{allocations[1]})); + if (cudaSuccess != error) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + error = CubDebug(detail::DebugSyncStream(stream)); + if (cudaSuccess != error) + { + return error; + } + + // Prepare streaming context for next partition (swap double buffers, advance number of processed items, etc.) + streaming_context.advance(current_num_items, (partition_idx + OffsetT{2} == num_partitions)); + } + } while (false); + + return error; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke() + { + return Invoke( + detail::scan::DeviceCompactInitKernel, + detail::select::DeviceSelectSweepKernel< + detail::select::policy_selector_from_hub, + InputIteratorT, + FlagsInputIteratorT, + SelectedOutputIteratorT, + NumSelectedIteratorT, + ScanTileStateT, + SelectOpT, + EqualityOpT, + per_partition_offset_t, + streaming_context_t, + SelectionOpt>); + } + + /** + * Internal dispatch routine + * + * @param 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 temp_storage_bytes + * Reference to size in bytes of `d_temp_storage` allocation + * + * @param d_in + * Pointer to the input sequence of data items + * + * @param d_flags + * Pointer to the input sequence of selection flags (if applicable) + * + * @param d_selected_out + * Pointer to the output sequence of selected data items + * + * @param d_num_selected_out + * Pointer to the total number of items selected (i.e., length of `d_selected_out`) + * + * @param select_op + * Selection operator + * + * @param equality_op + * Equality operator + * + * @param num_items + * Total number of input items (i.e., length of `d_in`) + * + * @param stream + * CUDA stream to launch kernels within. Default is stream0. + */ + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FlagsInputIteratorT d_flags, + SelectedOutputIteratorT d_selected_out, + NumSelectedIteratorT d_num_selected_out, + SelectOpT select_op, + EqualityOpT equality_op, + OffsetT num_items, + cudaStream_t stream) + { + int ptx_version = 0; + if (cudaError_t error = CubDebug(PtxVersion(ptx_version))) + { + return error; + } + + DispatchSelectIf dispatch( + d_temp_storage, + temp_storage_bytes, + d_in, + d_flags, + d_selected_out, + d_num_selected_out, + select_op, + equality_op, + num_items, + stream, + ptx_version); + + return CubDebug(PolicyHub::MaxPolicy::Invoke(ptx_version, dispatch)); + } +}; + +namespace detail::select +{ +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_policy( + [[maybe_unused]] PolicyGetter policy_getter, + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FlagsInputIteratorT d_flags, + SelectedOutputIteratorT d_selected_out, + NumSelectedIteratorT d_num_selected_out, + SelectOpT select_op, + EqualityOpT equality_op, + OffsetT num_items, + cudaStream_t stream, + [[maybe_unused]] PolicySelector policy_selector, + KernelLauncherFactory launcher_factory) +{ + static_assert(::cuda::std::is_empty_v); + + static constexpr bool is_partitioning_invocation = (SelectionOpt == SelectImpl::Partition); + static constexpr bool use_streaming_context = + (!is_partitioning_invocation) + || (static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max()) + < static_cast<::cuda::std::uint64_t>(::cuda::std::numeric_limits::max())); + using streaming_context_t = streaming_context_t; + using ScanTileStateT = ScanTileState; + static constexpr int init_kernel_threads = 128; + + using vsmem_helper_t = typename make_vsmem_helper< + PolicyGetter, + SelectionOpt, + InputIteratorT, + FlagsInputIteratorT, + SelectedOutputIteratorT, + SelectOpT, + EqualityOpT, + per_partition_offset_t, + streaming_context_t>::type; + + constexpr auto threads_per_block = vsmem_helper_t::agent_policy_t::BLOCK_THREADS; + constexpr auto items_per_thread = vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD; + constexpr auto tile_size = OffsetT{threads_per_block * items_per_thread}; + + static constexpr auto max_supported_partition_size = ::cuda::std::numeric_limits::max(); + static constexpr auto full_tile_partition_size = + max_supported_partition_size - (max_supported_partition_size % (threads_per_block * items_per_thread)); + static constexpr per_partition_offset_t capped_partition_size = + is_partitioning_invocation ? max_supported_partition_size : full_tile_partition_size; + + const auto max_partition_size = + (use_streaming_context && num_items > static_cast(capped_partition_size)) + ? static_cast(capped_partition_size) + : num_items; + const auto num_partitions = + (max_partition_size == 0) ? static_cast(1) : ::cuda::ceil_div(num_items, max_partition_size); + const auto max_num_tiles_per_invocation = static_cast(::cuda::ceil_div(max_partition_size, tile_size)); + const auto vsmem_size = max_num_tiles_per_invocation * vsmem_helper_t::vsmem_per_block; + + ::cuda::std::size_t streaming_selection_storage_bytes = + (num_partitions > 1) ? 2 * sizeof(OffsetT) : ::cuda::std::size_t{0}; + ::cuda::std::size_t allocation_sizes[3] = {0ULL, vsmem_size, streaming_selection_storage_bytes}; + + if (const auto error = + CubDebug(ScanTileStateT::AllocationSize(static_cast(max_num_tiles_per_invocation), allocation_sizes[0]))) + { + return error; + } + + void* allocations[3] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + OffsetT* tmp_num_selected_out = reinterpret_cast(allocations[2]); + streaming_context_t streaming_context{ + tmp_num_selected_out, (tmp_num_selected_out + 1), num_items, (num_partitions <= 1)}; + + for (OffsetT partition_idx = 0; partition_idx < num_partitions; partition_idx++) + { + OffsetT current_partition_offset = partition_idx * max_partition_size; + OffsetT current_num_items = + (partition_idx + 1 == num_partitions) ? (num_items - current_partition_offset) : max_partition_size; + + const auto current_num_tiles = static_cast(::cuda::ceil_div(current_num_items, tile_size)); + ScanTileStateT tile_status; + if (const auto error = CubDebug(tile_status.Init(current_num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + const int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(current_num_tiles, init_kernel_threads)); + +#ifdef CUB_DEBUG_LOG + _CubLog( + "Invoking scan_init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, init_kernel_threads, (long long) stream); +#endif + + if (const auto error = CubDebug( + launcher_factory(init_grid_size, init_kernel_threads, 0, stream) + .doit(detail::scan::DeviceCompactInitKernel, + tile_status, + current_num_tiles, + d_num_selected_out))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + if (current_num_items == 0) + { + return cudaSuccess; + } + +#ifdef CUB_DEBUG_LOG + { + int range_select_sm_occupancy; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + range_select_sm_occupancy, + DeviceSelectSweepKernel, + threads_per_block))) + { + return error; + } + + _CubLog("Invoking DeviceSelectSweepKernel<<<%d, %d, 0, " + "%lld>>>(), %d items per thread, %d SM occupancy\n", + current_num_tiles, + threads_per_block, + (long long) stream, + items_per_thread, + range_select_sm_occupancy); + } +#endif + + if (const auto error = CubDebug( + launcher_factory(current_num_tiles, threads_per_block, 0, stream) + .doit( + DeviceSelectSweepKernel, + d_in, + d_flags, + d_selected_out, + d_num_selected_out, + tile_status, + select_op, + equality_op, + static_cast(current_num_items), + current_num_tiles, + streaming_context, + cub::detail::vsmem_t{allocations[1]}))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + streaming_context.advance(current_num_items, (partition_idx + OffsetT{2} == num_partitions)); + } + + return cudaSuccess; +} + +template < + SelectImpl SelectionOpt, + typename InputIteratorT, + typename FlagsInputIteratorT, + typename SelectedOutputIteratorT, + typename NumSelectedIteratorT, + typename SelectOpT, + typename EqualityOpT, + typename OffsetT, + typename PolicySelector = + policy_selector_from_types, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires select_if_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FlagsInputIteratorT d_flags, + SelectedOutputIteratorT d_selected_out, + NumSelectedIteratorT d_num_selected_out, + SelectOpT select_op, + EqualityOpT equality_op, + OffsetT num_items, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelLauncherFactory launcher_factory = {}) +{ + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << PolicySelector{}(cc); + _CubLog("Dispatching DeviceSelectIf to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + return dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) { + return dispatch_policy( + policy_getter, + d_temp_storage, + temp_storage_bytes, + d_in, + d_flags, + d_selected_out, + d_num_selected_out, + select_op, + equality_op, + num_items, + stream, + policy_selector, + launcher_factory); + }); +} +} // namespace detail::select + +CUB_NAMESPACE_END + +_CCCL_DIAG_POP diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_streaming_reduce.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_streaming_reduce.cuh new file mode 100644 index 00000000..ceb8391b --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_streaming_reduce.cuh @@ -0,0 +1,323 @@ +// SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +#ifndef _CCCL_DOXYGEN_INVOKED // Do not document + +CUB_NAMESPACE_BEGIN + +namespace detail::reduce +{ +template +struct accumulating_transform_output_op +{ + bool first_partition; + bool last_partition; + + // We use a double-buffer to make assignment idempotent (i.e., allow potential repeated assignment) + GlobalAccumT* d_previous_aggregate; + GlobalAccumT* d_aggregate_out; + + // Output iterator to which the final result of type `GlobalAccumT` across all partitions will be assigned + FinalResultOutIteratorT d_out; + + // Unary promotion operator type that is used to transform a per-partition result to a global result + PromoteToGlobalOpT promote_op; + + // Reduction operation + GlobalReductionOpT reduce_op; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(IndexT, AccumT per_partition_aggregate) + { + // Add this partitions aggregate to the global aggregate + if (first_partition) + { + *d_aggregate_out = promote_op(per_partition_aggregate); + } + else + { + *d_aggregate_out = reduce_op(*d_previous_aggregate, promote_op(per_partition_aggregate)); + } + + // If this is the last partition, we write the global aggregate to the user-provided iterator + if (last_partition) + { + *d_out = *d_aggregate_out; + } + } + + /** + * This is a helper function that's invoked after a partition has been fully processed + */ + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(GlobalOffsetT partition_size, bool next_partition_is_the_last) + { + promote_op.advance(partition_size); + using ::cuda::std::swap; + swap(d_previous_aggregate, d_aggregate_out); + first_partition = false; + last_partition = next_partition_is_the_last; + } +}; + +/** + * Unary "promotion" operator type that is used to transform a per-partition result to a global result + */ +template +struct local_to_global_op +{ + // The current partition's offset to be factored into this partition's index + GlobalOffsetT current_partition_offset; + + /** + * This helper function is invoked after a partition has been fully processed, in preparation for the next partition. + */ + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE void advance(GlobalOffsetT partition_size) + { + current_partition_offset += partition_size; + } + + /** + * Unary operator called to transform the per-partition aggregate of a partition to a global aggregate type (i.e., one + * that is used to reduce across partitions). + */ + template + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE KeyValuePair + operator()(KeyValuePair partition_aggregate) + { + return KeyValuePair{ + current_partition_offset + static_cast(partition_aggregate.key), partition_aggregate.value}; + } +}; + +template +struct unzip_and_write_arg_extremum_op +{ + ExtremumOutIteratorT result_out_it; + IndexOutIteratorT index_out_it; + + template + _CCCL_DEVICE _CCCL_FORCEINLINE void operator()(IndexT, KeyValuePairT reduced_result) + { + *result_out_it = reduced_result.value; + *index_out_it = reduced_result.key; + } +}; + +/****************************************************************************** + * Single-problem streaming reduction dispatch + *****************************************************************************/ + +// Internal dispatch routine for computing a device-wide argument extremum, like `ArgMin` and `ArgMax`. +// Streaming, here, refers to the approach used for large number of items that are processed in multiple partitions. +// +// @tparam PerPartitionOffsetT +// Offset type used as the index to access items within one partition, i.e., the offset type used within the kernel +// template specialization +// +// @tparam InputIteratorT +// Random-access input iterator type for reading input items @iterator +// +// @tparam OutputIteratorT +// Output iterator type for writing the result of the (index, extremum)-key-value-pair +// +// @tparam GlobalOffsetT +// Offset type used as the index to access items within the total input range, i.e., in the range [d_in, d_in + +// num_items) +// +// @tparam ReductionOpT +// Binary reduction functor type having a member function that returns the selected extremum of two input items. +// The streaming reduction requires two overloads, one used for selecting the extremum within one partition and one +// for selecting the extremum across partitions. +// +// @tparam TuningEnvT +// Tuning environment Environment type +// +template +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_streaming_arg_reduce( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + ExtremumOutIteratorT d_min_out, + IndexOutIteratorT d_index_out, + GlobalOffsetT num_items, + ReductionOpT reduce_op, + cudaStream_t stream, + const TuningEnvT& = {}) +{ + using input_value_t = detail::it_value_t; + using output_extremum_t = detail::non_void_value_t; + using default_policy_selector_t = detail::reduce:: + policy_selector_from_types, PerPartitionOffsetT, ReductionOpT>; + using default_policy_t = decltype(default_policy_selector_t{}(::cuda::compute_capability{})); + using policy_selector_t = + ::cuda::std::execution::__query_result_or_t; + +# if _CCCL_HAS_CONCEPTS() + static_assert(reduce_policy_selector); +# endif // _CCCL_HAS_CONCEPTS() + + // Tabulate output iterator that unzips the result and writes it to the user-provided output iterators + auto d_result_out = ::cuda::make_tabulate_output_iterator( + detail::reduce::unzip_and_write_arg_extremum_op{d_min_out, d_index_out}); + + // Wrapped input iterator to produce index-value tuples, i.e., -tuples + // We make sure to offset the user-provided input iterator by the current partition's offset + using arg_index_input_iterator_t = ArgIndexInputIterator; + + // The output tuple type (i.e., extremum plus index tuples) + using per_partition_accum_t = KeyValuePair; + using global_accum_t = KeyValuePair; + + // Unary promotion operator type that is used to transform a per-partition result to a global result + // operator()(per_partition_accum_t) -> global_accum_t + using local_to_global_op_t = local_to_global_op; + + // The current partition's input iterator is an ArgIndex iterator that generates indices relative to the beginning + // of the current partition, i.e., [0, partition_size) along with an OffsetIterator that offsets the user-provided + // input iterator by the current partition's offset + arg_index_input_iterator_t d_indexed_offset_in(d_in); + + // Transforms the per-partition result to a global result by adding the current partition's offset to the arg result + // of a partition + local_to_global_op_t local_to_global_op{GlobalOffsetT{0}}; + + // Upper bound at which we want to cut the input into multiple partitions. Align to 4096 bytes for performance + // reasons + static constexpr PerPartitionOffsetT max_offset_size = ::cuda::std::numeric_limits::max(); + static constexpr PerPartitionOffsetT max_partition_size = + max_offset_size - (max_offset_size % PerPartitionOffsetT{4096}); + + // Whether the given number of items fits into a single partition + const bool is_single_partition = + static_cast(max_partition_size) >= static_cast(num_items); + + // The largest partition size ever encountered + const auto largest_partition_size = + is_single_partition ? static_cast(num_items) : max_partition_size; + + // Reduction operator type that enables accumulating per-partition results to a global reduction result + using accumulating_transform_output_op_t = + accumulating_transform_output_op; + auto accumulating_out_op = accumulating_transform_output_op_t{ + true, is_single_partition, nullptr, nullptr, d_result_out, local_to_global_op, reduce_op}; + + // Initial value for empty problems, according to documented contract + const auto empty_problem_extremum = static_cast([] { + if constexpr (::cuda::std::is_same_v + && ::cuda::std::numeric_limits::is_specialized) + { + return ::cuda::std::numeric_limits::max(); + } + else if constexpr (::cuda::std::is_same_v + && ::cuda::std::numeric_limits::is_specialized) + { + return ::cuda::std::numeric_limits::lowest(); + } + else + { + return input_value_t{}; + } + }()); + auto initial_value = empty_problem_init_t{{PerPartitionOffsetT{1}, empty_problem_extremum}}; + + void* allocations[2] = {nullptr, nullptr}; + size_t allocation_sizes[2] = {0, 2 * sizeof(global_accum_t)}; + + // Query temporary storage requirements for per-partition reduction + reduce::dispatch( + nullptr, + allocation_sizes[0], + d_indexed_offset_in, + ::cuda::make_tabulate_output_iterator(accumulating_out_op), + static_cast(largest_partition_size), + reduce_op, + initial_value, + stream, + ::cuda::std::identity{}, + policy_selector_t{}); + + // Alias the temporary allocations from the single storage blob (or compute the necessary size of the blob) + if (const auto error = CubDebug(alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + // Return if the caller is simply requesting the size of the storage allocation + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + // Pointer to the double-buffer of global accumulators, which aggregate cross-partition results + global_accum_t* const d_global_aggregates = static_cast(allocations[1]); + + accumulating_out_op.d_previous_aggregate = d_global_aggregates; + accumulating_out_op.d_aggregate_out = d_global_aggregates + 1; + + for (GlobalOffsetT current_partition_offset = 0; current_partition_offset < static_cast(num_items); + current_partition_offset += static_cast(max_partition_size)) + { + const GlobalOffsetT remaining_items = (num_items - current_partition_offset); + const GlobalOffsetT current_num_items = + (remaining_items < max_partition_size) ? remaining_items : max_partition_size; + + d_indexed_offset_in = arg_index_input_iterator_t(d_in + current_partition_offset); + + if (const auto error = reduce::dispatch( + d_temp_storage, + temp_storage_bytes, + d_indexed_offset_in, + ::cuda::make_tabulate_output_iterator(accumulating_out_op), + static_cast(current_num_items), + reduce_op, + initial_value, + stream, + ::cuda::std::identity{}, + policy_selector_t{})) + { + return error; + } + + // Whether the next partition will be the last partition + const bool next_partition_is_last = + (remaining_items - current_num_items) <= static_cast(max_partition_size); + accumulating_out_op.advance(current_num_items, next_partition_is_last); + } + + return cudaSuccess; +} +} // namespace detail::reduce +CUB_NAMESPACE_END + +#endif // !_CCCL_DOXYGEN_INVOKED diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh new file mode 100644 index 00000000..f2395c0f --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_streaming_reduce_by_key.cuh @@ -0,0 +1,255 @@ +// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3-Clause + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::reduce_by_key +{ +template >, + typename PolicySelector = + policy_selector_from_types>>> +#if _CCCL_HAS_CONCEPTS() + requires reduce_by_key_policy_selector +#endif +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t dispatch_streaming( + void* d_temp_storage, + size_t& temp_storage_bytes, + KeysInputIteratorT d_keys_in, + UniqueOutputIteratorT d_unique_out, + ValuesInputIteratorT d_values_in, + AggregatesOutputIteratorT d_aggregates_out, + NumRunsOutputIteratorT d_num_runs_out, + EqualityOpT equality_op, + ReductionOpT reduction_op, + OffsetT num_items, + cudaStream_t stream, + PolicySelector policy_selector = {}) +{ + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(ptx_compute_cap(cc))) + { + return error; + } + + const ReduceByKeyPolicy policy = policy_selector(cc); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << policy; + _CubLog("Dispatching streaming reduce by key to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + using local_offset_t = ::cuda::std::int32_t; + using global_offset_t = OffsetT; + static constexpr bool use_streaming_invocation = + ::cuda::std::numeric_limits::max() > ::cuda::std::numeric_limits::max(); + using streaming_context_t = ::cuda::std:: + conditional_t, NullType>; + using ScanTileStateT = ReduceByKeyScanTileState; + [[maybe_unused]] static constexpr int init_kernel_threads = 128; + + const int threads_per_block = policy.lookback.threads_per_block; + const int items_per_thread = policy.lookback.items_per_thread; + const auto tile_size = + static_cast(threads_per_block) * static_cast(items_per_thread); + + auto capped_num_items_per_invocation = num_items; + if constexpr (use_streaming_invocation) + { + capped_num_items_per_invocation = static_cast(::cuda::std::numeric_limits::max()); + capped_num_items_per_invocation -= (capped_num_items_per_invocation % tile_size); + } + + const auto max_num_items_per_invocation = + use_streaming_invocation ? ::cuda::std::min(capped_num_items_per_invocation, num_items) : num_items; + auto const num_partitions = + (num_items == 0) ? global_offset_t{1} : ::cuda::ceil_div(num_items, capped_num_items_per_invocation); + + const auto max_num_tiles = static_cast(::cuda::ceil_div(max_num_items_per_invocation, tile_size)); + + size_t allocation_sizes[3]; + if (const auto error = CubDebug(ScanTileStateT::AllocationSize(max_num_tiles, allocation_sizes[0]))) + { + return error; + } + allocation_sizes[1] = num_partitions > 1 ? sizeof(global_offset_t) * 2 : size_t{0}; + allocation_sizes[2] = num_partitions > 1 ? sizeof(AccumT) * 2 : size_t{0}; + + void* allocations[3] = {}; + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + return cudaSuccess; + } + + for (global_offset_t partition_idx = 0; partition_idx < num_partitions; partition_idx++) + { + global_offset_t current_partition_offset = partition_idx * capped_num_items_per_invocation; + global_offset_t current_num_items = + (partition_idx + 1 == num_partitions) ? (num_items - current_partition_offset) : capped_num_items_per_invocation; + + const auto num_current_tiles = static_cast(::cuda::ceil_div(current_num_items, tile_size)); + ScanTileStateT tile_state; + if (const auto error = CubDebug(tile_state.Init(num_current_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + const int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(num_current_tiles, init_kernel_threads)); +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, init_kernel_threads, (long long) stream); +#endif + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(init_grid_size, init_kernel_threads, 0, stream) + .doit(&detail::scan::DeviceCompactInitKernel, + tile_state, + num_current_tiles, + d_num_runs_out))) + { + return error; + } + + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + if (num_items == 0) + { + return cudaSuccess; + } + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking reduce_by_key_kernel<<<%d, %d, 0, %lld>>>(), %d items per thread\n", + num_current_tiles, + threads_per_block, + (long long) stream, + items_per_thread); +#endif + auto reduce_by_key_kernel = DeviceReduceByKeyKernel< + PolicySelector, + KeysInputIteratorT, + UniqueOutputIteratorT, + ValuesInputIteratorT, + AggregatesOutputIteratorT, + NumRunsOutputIteratorT, + ScanTileStateT, + EqualityOpT, + ReductionOpT, + local_offset_t, + AccumT, + streaming_context_t>; + + if constexpr (use_streaming_invocation) + { + auto tmp_num_uniques = static_cast(allocations[1]); + auto tmp_prefix = static_cast(allocations[2]); + const bool is_first_partition = (partition_idx == 0); + const bool is_last_partition = (partition_idx + 1 == num_partitions); + const int buffer_selector = partition_idx % 2; + streaming_context_t streaming_context{ + is_first_partition, + is_last_partition, + is_first_partition ? d_keys_in : d_keys_in + current_partition_offset - 1, + &tmp_prefix[buffer_selector], + &tmp_prefix[buffer_selector ^ 0x01], + &tmp_num_uniques[buffer_selector], + &tmp_num_uniques[buffer_selector ^ 0x01]}; + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_current_tiles, threads_per_block, 0, stream) + .doit(reduce_by_key_kernel, + d_keys_in + current_partition_offset, + d_unique_out, + d_values_in + current_partition_offset, + d_aggregates_out, + d_num_runs_out, + tile_state, + 0, + equality_op, + reduction_op, + static_cast(current_num_items), + streaming_context, + detail::vsmem_t{nullptr}))) + { + return error; + } + } + else + { + if (const auto error = CubDebug( + THRUST_NS_QUALIFIER::cuda_cub::detail::triple_chevron(num_current_tiles, threads_per_block, 0, stream) + .doit(reduce_by_key_kernel, + d_keys_in + current_partition_offset, + d_unique_out, + d_values_in + current_partition_offset, + d_aggregates_out, + d_num_runs_out, + tile_state, + 0, + equality_op, + reduction_op, + static_cast(current_num_items), + NullType{}, + detail::vsmem_t{nullptr}))) + { + return error; + } + } + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + } + + return cudaSuccess; +} +} // namespace detail::reduce_by_key + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_three_way_partition.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_three_way_partition.cuh new file mode 100644 index 00000000..13ab3b78 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_three_way_partition.cuh @@ -0,0 +1,670 @@ +// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved. +// SPDX-License-Identifier: BSD-3 + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::three_way_partition +{ +template +struct DeviceThreeWayPartitionKernelSource +{ + CUB_DEFINE_KERNEL_GETTER(ThreeWayPartitionInitKernel, + DeviceThreeWayPartitionInitKernel); + + CUB_DEFINE_KERNEL_GETTER( + ThreeWayPartitionKernel, + DeviceThreeWayPartitionKernel< + PolicySelector, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + per_partition_offset_t, + streaming_context_t>); +}; + +// TODO(bgruber): remove in CCCL 4.0 +template +struct policy_selector_from_hub +{ + [[nodiscard]] _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability /*cc*/) const + -> ThreeWayPartitionPolicy + { + using active_policy = typename PolicyHub::MaxPolicy::ActivePolicy::ThreeWayPartitionPolicy; + return ThreeWayPartitionPolicy{ + ThreeWayPartitionAlgorithm::lookback, + {active_policy::BLOCK_THREADS, + active_policy::ITEMS_PER_THREAD, + active_policy::LOAD_ALGORITHM, + active_policy::LOAD_MODIFIER, + active_policy::SCAN_ALGORITHM, + lookback_delay_policy_from_type}}; + } +}; + +// TODO(bgruber): drop in CCCL 4.0 when we drop the segmented sort dispatcher, which depends on this +template < + typename InputIteratorT, + typename FirstOutputIteratorT, + typename SecondOutputIteratorT, + typename UnselectedOutputIteratorT, + typename NumSelectedIteratorT, + typename SelectFirstPartOp, + typename SelectSecondPartOp, + typename OffsetT, + typename PolicyHub = detail::three_way_partition::policy_hub, + detail::three_way_partition::per_partition_offset_t>, + typename KernelSource = detail::three_way_partition::DeviceThreeWayPartitionKernelSource< + detail::three_way_partition::policy_selector_from_hub, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + detail::three_way_partition::ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + detail::three_way_partition::per_partition_offset_t, + detail::three_way_partition::streaming_context_t, + OffsetT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +struct dispatch_three_way_partition_if +{ + /***************************************************************************** + * Types and constants + ****************************************************************************/ + + // Offset type used to instantiate the three-way partition-kernel and agent to index the items within one partition + using per_partition_offset_t = detail::three_way_partition::per_partition_offset_t; + + // Type used to provide streaming information about each partition's context + static constexpr per_partition_offset_t partition_size = ::cuda::std::numeric_limits::max(); + + using streaming_context_t = detail::three_way_partition::streaming_context_t; + + using ScanTileStateT = detail::three_way_partition::ScanTileStateT; + + static constexpr int INIT_KERNEL_THREADS = 256; + + void* d_temp_storage; + size_t& temp_storage_bytes; + InputIteratorT d_in; + FirstOutputIteratorT d_first_part_out; + SecondOutputIteratorT d_second_part_out; + UnselectedOutputIteratorT d_unselected_out; + NumSelectedIteratorT d_num_selected_out; + SelectFirstPartOp select_first_part_op; + SelectSecondPartOp select_second_part_op; + OffsetT num_items; + cudaStream_t stream; + KernelSource kernel_source; + KernelLauncherFactory launcher_factory; + + /***************************************************************************** + * Dispatch entrypoints + ****************************************************************************/ + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t __invoke( + int threads_per_block, + int items_per_thread, + ScanInitKernelPtrT three_way_partition_init_kernel, + SelectIfKernelPtrT three_way_partition_kernel) + { + const int tile_size = threads_per_block * items_per_thread; + + // The maximum number of items for which we will ever invoke the kernel (i.e. largest partition size) + auto const max_partition_size = static_cast( + (::cuda::std::min) (static_cast(num_items), static_cast(partition_size))); + + // The number of partitions required to "iterate" over the total input + auto const num_partitions = + (max_partition_size == 0) ? OffsetT{1} : ::cuda::ceil_div(num_items, max_partition_size); + + // The maximum number of tiles for which we will ever invoke the kernel + auto const max_num_tiles_per_invocation = static_cast(::cuda::ceil_div(max_partition_size, tile_size)); + + // For streaming invocations, we need two sets (for double-buffering) of three counters each + constexpr ::cuda::std::size_t num_counters_per_pass = 3; + constexpr ::cuda::std::size_t num_streaming_counters = 2 * num_counters_per_pass; + ::cuda::std::size_t streaming_selection_storage_bytes = + (num_partitions > 1) ? num_streaming_counters * sizeof(OffsetT) : ::cuda::std::size_t{0}; + + // Specify temporary storage allocation requirements + size_t allocation_sizes[2] = {0ULL, streaming_selection_storage_bytes}; + + if (const auto error = + CubDebug(ScanTileStateT::AllocationSize(static_cast(max_num_tiles_per_invocation), allocation_sizes[0]))) + { + return error; + } + + // Compute allocation pointers into the single storage blob (or compute the necessary size of the blob) + void* allocations[2] = {}; + + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage + // allocation + return cudaSuccess; + } + + // Initialize the streaming context with the temporary storage for double-buffering the previously selected items + // and the total number (across all partitions) of items + OffsetT* tmp_num_selected_out = static_cast(allocations[1]); + streaming_context_t streaming_context{ + tmp_num_selected_out, (tmp_num_selected_out + num_counters_per_pass), (num_partitions <= 1)}; + + // Iterate over the partitions until all input is processed + for (OffsetT partition_idx = 0; partition_idx < num_partitions; partition_idx++) + { + OffsetT current_partition_offset = partition_idx * max_partition_size; + OffsetT current_num_items = + (partition_idx + 1 == num_partitions) ? (num_items - current_partition_offset) : max_partition_size; + + // Construct the tile status interface + const auto current_num_tiles = static_cast(::cuda::ceil_div(current_num_items, tile_size)); + + // Construct the tile status interface + ScanTileStateT tile_status; + if (const auto error = CubDebug(tile_status.Init(current_num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + // Log three_way_partition_init_kernel configuration + const int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(current_num_tiles, INIT_KERNEL_THREADS)); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking three_way_partition_init_kernel<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + INIT_KERNEL_THREADS, + reinterpret_cast(stream)); +#endif // CUB_DEBUG_LOG + + // Invoke three_way_partition_init_kernel to initialize tile descriptors + if (const auto error = CubDebug( + launcher_factory(init_grid_size, INIT_KERNEL_THREADS, 0, stream) + .doit(three_way_partition_init_kernel, tile_status, current_num_tiles, d_num_selected_out))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // No more items to process (note, we do not want to return early for num_items==0, because we need to make sure + // that `three_way_partition_init_kernel` has written '0' to d_num_selected_out) + if (current_num_items == 0) + { + return cudaSuccess; + } + +// Log select_if_kernel configuration +#ifdef CUB_DEBUG_LOG + { + // Get SM occupancy for select_if_kernel + int range_select_sm_occupancy; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + range_select_sm_occupancy, // out + three_way_partition_kernel, + threads_per_block))) + { + return error; + } + + _CubLog("Invoking three_way_partition_kernel<<<%d, %d, 0, %lld>>>(), %d " + "items per thread, %d SM occupancy\n", + current_num_tiles, + threads_per_block, + reinterpret_cast(stream), + items_per_thread, + range_select_sm_occupancy); + } +#endif // CUB_DEBUG_LOG + + // Invoke select_if_kernel + if (const auto error = CubDebug( + launcher_factory(current_num_tiles, threads_per_block, 0, stream) + .doit(three_way_partition_kernel, + d_in, + d_first_part_out, + d_second_part_out, + d_unselected_out, + d_num_selected_out, + tile_status, + select_first_part_op, + select_second_part_op, + static_cast(current_num_items), + current_num_tiles, + streaming_context))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Prepare streaming context for next partition (swap double buffers, advance number of processed items, etc.) + streaming_context.advance(current_num_items, (partition_idx + OffsetT{2} == num_partitions)); + } + + return cudaSuccess; + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t + Invoke(ActivePolicyT policy, + ScanInitKernelPtrT three_way_partition_init_kernel, + SelectIfKernelPtrT three_way_partition_kernel) + { + const int threads_per_block = policy.ThreeWayPartition().ThreadsPerBlock(); + const int items_per_thread = policy.ThreeWayPartition().ItemsPerThread(); + return __invoke(threads_per_block, items_per_thread, three_way_partition_init_kernel, three_way_partition_kernel); + } + + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT active_policy = {}) + { + const auto wrapped_policy = detail::three_way_partition::MakeThreeWayPartitionPolicyWrapper(active_policy); + return Invoke(wrapped_policy, kernel_source.ThreeWayPartitionInitKernel(), kernel_source.ThreeWayPartitionKernel()); + } + + /** + * Internal dispatch routine + */ + template + CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FirstOutputIteratorT d_first_part_out, + SecondOutputIteratorT d_second_part_out, + UnselectedOutputIteratorT d_unselected_out, + NumSelectedIteratorT d_num_selected_out, + SelectFirstPartOp select_first_part_op, + SelectSecondPartOp select_second_part_op, + OffsetT num_items, + cudaStream_t stream, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}, + MaxPolicyT max_policy = {}) + { + // Get PTX version + int ptx_version = 0; + if (cudaError error = CubDebug(launcher_factory.PtxVersion(ptx_version)); cudaSuccess != error) + { + return error; + } + + dispatch_three_way_partition_if dispatch{ + d_temp_storage, + temp_storage_bytes, + d_in, + d_first_part_out, + d_second_part_out, + d_unselected_out, + d_num_selected_out, + select_first_part_op, + select_second_part_op, + num_items, + stream, + kernel_source, + launcher_factory}; + + return CubDebug(max_policy.Invoke(ptx_version, dispatch)); + } +}; +} // namespace detail::three_way_partition + +/****************************************************************************** + * Dispatch + ******************************************************************************/ + +// TODO(bgruber): Drop in CCCL 4.0 +//! Deprecated [Since 3.5] +template < + typename InputIteratorT, + typename FirstOutputIteratorT, + typename SecondOutputIteratorT, + typename UnselectedOutputIteratorT, + typename NumSelectedIteratorT, + typename SelectFirstPartOp, + typename SelectSecondPartOp, + typename OffsetT, + typename PolicyHub = detail::three_way_partition::policy_hub, + detail::three_way_partition::per_partition_offset_t>, + typename KernelSource = detail::three_way_partition::DeviceThreeWayPartitionKernelSource< + detail::three_way_partition::policy_selector_from_hub, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + detail::three_way_partition::ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + detail::three_way_partition::per_partition_offset_t, + detail::three_way_partition::streaming_context_t, + OffsetT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +using DispatchThreeWayPartitionIf CCCL_DEPRECATED_BECAUSE("Use the tuning API for DevicePartition") = + detail::three_way_partition::dispatch_three_way_partition_if< + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + SelectFirstPartOp, + SelectSecondPartOp, + OffsetT, + PolicyHub, + KernelSource, + KernelLauncherFactory>; + +namespace detail::three_way_partition +{ +template , per_partition_offset_t>, + typename KernelSource = DeviceThreeWayPartitionKernelSource< + PolicySelector, + InputIteratorT, + FirstOutputIteratorT, + SecondOutputIteratorT, + UnselectedOutputIteratorT, + NumSelectedIteratorT, + ScanTileStateT, + SelectFirstPartOp, + SelectSecondPartOp, + per_partition_offset_t, + streaming_context_t, + OffsetT>, + typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY> +#if _CCCL_HAS_CONCEPTS() + requires three_way_partition_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch( + void* d_temp_storage, + size_t& temp_storage_bytes, + InputIteratorT d_in, + FirstOutputIteratorT d_first_part_out, + SecondOutputIteratorT d_second_part_out, + UnselectedOutputIteratorT d_unselected_out, + NumSelectedIteratorT d_num_selected_out, + SelectFirstPartOp select_first_part_op, + SelectSecondPartOp select_second_part_op, + OffsetT num_items, + cudaStream_t stream, + PolicySelector policy_selector = {}, + KernelSource kernel_source = {}, + KernelLauncherFactory launcher_factory = {}) +{ + ::cuda::compute_capability cc{}; + if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc))) + { + return error; + } + + const ThreeWayPartitionPolicy active_policy = policy_selector(cc); + +#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + NV_IF_TARGET(NV_IS_HOST, ({ + ::std::stringstream ss; + ss << active_policy; + _CubLog("Dispatching DeviceThreeWayPartition to compute capability %d.%d with tuning: %s\n", + cc.major_cap(), + cc.minor_cap(), + ss.str().c_str()); + })) +#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG) + + static constexpr per_partition_offset_t partition_size = ::cuda::std::numeric_limits::max(); + static constexpr int init_kernel_threads = 256; + + const int threads_per_block = active_policy.lookback.threads_per_block; + const int items_per_thread = active_policy.lookback.items_per_thread; + const int tile_size = threads_per_block * items_per_thread; + + auto three_way_partition_init_kernel = kernel_source.ThreeWayPartitionInitKernel(); + auto three_way_partition_kernel = kernel_source.ThreeWayPartitionKernel(); + + // The maximum number of items for which we will ever invoke the kernel (i.e. largest partition size) + auto const max_partition_size = + static_cast((::cuda::std::min) (static_cast(num_items), static_cast(partition_size))); + + // The number of partitions required to "iterate" over the total input + auto const num_partitions = (max_partition_size == 0) ? OffsetT{1} : ::cuda::ceil_div(num_items, max_partition_size); + + // The maximum number of tiles for which we will ever invoke the kernel + auto const max_num_tiles_per_invocation = static_cast(::cuda::ceil_div(max_partition_size, tile_size)); + + // For streaming invocations, we need two sets (for double-buffering) of three counters each + constexpr ::cuda::std::size_t num_counters_per_pass = 3; + constexpr ::cuda::std::size_t num_streaming_counters = 2 * num_counters_per_pass; + ::cuda::std::size_t streaming_selection_storage_bytes = + (num_partitions > 1) ? num_streaming_counters * sizeof(OffsetT) : ::cuda::std::size_t{0}; + + // Specify temporary storage allocation requirements + size_t allocation_sizes[2] = {0ULL, streaming_selection_storage_bytes}; + + if (const auto error = + CubDebug(ScanTileStateT::AllocationSize(static_cast(max_num_tiles_per_invocation), allocation_sizes[0]))) + { + return error; + } + + // Compute allocation pointers into the single storage blob (or compute the necessary size of the blob) + void* allocations[2] = {}; + + if (const auto error = + CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes))) + { + return error; + } + + if (d_temp_storage == nullptr) + { + // Return if the caller is simply requesting the size of the storage allocation + return cudaSuccess; + } + + // Initialize the streaming context with the temporary storage for double-buffering the previously selected items + // and the total number (across all partitions) of items + OffsetT* tmp_num_selected_out = static_cast(allocations[1]); + streaming_context_t streaming_context{ + tmp_num_selected_out, (tmp_num_selected_out + num_counters_per_pass), (num_partitions <= 1)}; + + // Iterate over the partitions until all input is processed + for (OffsetT partition_idx = 0; partition_idx < num_partitions; partition_idx++) + { + OffsetT current_partition_offset = partition_idx * max_partition_size; + OffsetT current_num_items = + (partition_idx + 1 == num_partitions) ? (num_items - current_partition_offset) : max_partition_size; + + // Construct the tile status interface + const auto current_num_tiles = static_cast(::cuda::ceil_div(current_num_items, tile_size)); + + // Construct the tile status interface + ScanTileStateT tile_status; + if (const auto error = CubDebug(tile_status.Init(current_num_tiles, allocations[0], allocation_sizes[0]))) + { + return error; + } + + // Log three_way_partition_init_kernel configuration + const int init_grid_size = ::cuda::std::max(1, ::cuda::ceil_div(current_num_tiles, init_kernel_threads)); + +#ifdef CUB_DEBUG_LOG + _CubLog("Invoking three_way_partition_init_kernel<<<%d, %d, 0, %lld>>>()\n", + init_grid_size, + init_kernel_threads, + reinterpret_cast(stream)); +#endif // CUB_DEBUG_LOG + + // Invoke three_way_partition_init_kernel to initialize tile descriptors + if (const auto error = CubDebug( + launcher_factory(init_grid_size, init_kernel_threads, 0, stream) + .doit(three_way_partition_init_kernel, tile_status, current_num_tiles, d_num_selected_out))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // No more items to process (note, we do not want to return early for num_items==0, because we need to make sure + // that `three_way_partition_init_kernel` has written '0' to d_num_selected_out) + if (current_num_items == 0) + { + return cudaSuccess; + } + + // Log select_if_kernel configuration +#ifdef CUB_DEBUG_LOG + { + // Get SM occupancy for three_way_partition_kernel + int range_select_sm_occupancy; + if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy( + range_select_sm_occupancy, // out + three_way_partition_kernel, + threads_per_block))) + { + return error; + } + + _CubLog("Invoking three_way_partition_kernel<<<%d, %d, 0, %lld>>>(), %d " + "items per thread, %d SM occupancy\n", + current_num_tiles, + threads_per_block, + reinterpret_cast(stream), + items_per_thread, + range_select_sm_occupancy); + } +#endif // CUB_DEBUG_LOG + + // Invoke three_way_partition_kernel + if (const auto error = CubDebug( + launcher_factory(current_num_tiles, threads_per_block, 0, stream) + .doit(three_way_partition_kernel, + d_in, + d_first_part_out, + d_second_part_out, + d_unselected_out, + d_num_selected_out, + tile_status, + select_first_part_op, + select_second_part_op, + static_cast(current_num_items), + current_num_tiles, + streaming_context))) + { + return error; + } + + // Check for failure to launch + if (const auto error = CubDebug(cudaPeekAtLastError())) + { + return error; + } + + // Sync the stream if specified to flush runtime errors + if (const auto error = CubDebug(detail::DebugSyncStream(stream))) + { + return error; + } + + // Prepare streaming context for next partition (swap double buffers, advance number of processed items, etc.) + streaming_context.advance(current_num_items, (partition_idx + OffsetT{2} == num_partitions)); + } + + return cudaSuccess; +} +} // namespace detail::three_way_partition + +CUB_NAMESPACE_END diff --git a/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_topk.cuh b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_topk.cuh new file mode 100644 index 00000000..18adf889 --- /dev/null +++ b/qwen3_6_scripts/cccl_preload/include/cub/device/dispatch/dispatch_topk.cuh @@ -0,0 +1,749 @@ +// 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 unordered data items residing within device-accessible memory. + +#pragma once + +#include + +#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) +# pragma GCC system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) +# pragma clang system_header +#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) +# pragma system_header +#endif // no system header + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +CUB_NAMESPACE_BEGIN + +namespace detail::topk +{ +// Used in the bin ID calculation to exclude bits unrelated to the current pass +template +[[nodiscard]] _CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr unsigned calc_mask(const int pass) +{ + int num_bits = calc_start_bit(pass - 1) - calc_start_bit(pass); + return (1 << num_bits) - 1; +} + +// Get the bin ID from the value of element +template > +struct extract_bin_op_t; + +template +struct extract_bin_op_t +{ + static constexpr bool is_descending = SelectDirection != select::min; + using bit_ordered_type = typename Traits::UnsignedBits; + + int pass{}; + int start_bit{}; + unsigned mask{}; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE extract_bin_op_t(int pass, int /*total_bits*/, DecomposerT /*decomposer*/) + : pass(pass) + , start_bit(calc_start_bit(pass)) + , mask(calc_mask(pass)) + {} + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE int operator()(T key) const + { + auto bits = reinterpret_cast::UnsignedBits&>(key); + bits = Traits::TwiddleIn(bits); + if constexpr (SelectDirection != select::min) + { + bits = ~bits; + } + int bucket = (bits >> start_bit) & mask; + return bucket; + } +}; + +template +struct extract_bin_op_t +{ + static constexpr bool is_descending = SelectDirection != select::min; + using radix_traits_t = detail::radix::traits_t; + using bit_ordered_type = typename radix_traits_t::bit_ordered_type; + using digit_extractor_t = typename radix_traits_t::template digit_extractor_t, DecomposerT>; + + DecomposerT decomposer{}; + digit_extractor_t digit_extractor; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE extract_bin_op_t(int pass, int total_bits, DecomposerT decomposer) + : decomposer(decomposer) + , digit_extractor(radix_traits_t::template digit_extractor>( + calc_start_bit(total_bits, pass), + calc_start_bit(total_bits, pass - 1) - calc_start_bit(total_bits, pass), + decomposer)) + {} + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE int operator()(T key) const + { + bit_ordered_type ordered = key; + ordered = RadixSortTwiddle::In(ordered, decomposer); + return static_cast(digit_extractor.Digit(ordered)); + } +}; + +// Check if the input element is still a candidate for the target pass. +template > +struct identify_candidates_op_t; + +template +struct identify_candidates_op_t +{ + using unsigned_bits_t = typename Traits::UnsignedBits; + using key_prefix_t = key_prefix_storage_t; + unsigned_bits_t* kth_key_bits; + int start_bit; + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE + identify_candidates_op_t(key_prefix_t* kth_key_bits, int pass, int /*total_bits*/, DecomposerT /*decomposer*/) + : kth_key_bits(&kth_key_bits->bits) + { + start_bit = calc_start_bit(pass - 1); + } + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE candidate_class operator()(T key) const + { + auto bits = reinterpret_cast(key); + bits = Traits::TwiddleIn(bits); + + if constexpr (SelectDirection != select::min) + { + bits = ~bits; + } + + bits = (bits >> start_bit) << start_bit; + + return (bits < *kth_key_bits) ? candidate_class::selected + : (bits == *kth_key_bits) + ? candidate_class::candidate + : candidate_class::rejected; + } +}; + +template +struct identify_candidates_op_t +{ + static constexpr bool is_descending = SelectDirection != select::min; + using radix_traits_t = detail::radix::traits_t; + using bit_ordered_type = typename radix_traits_t::bit_ordered_type; + using key_prefix_t = key_prefix_storage_t; + + key_prefix_t* kth_key_bits{}; + int pass{}; + int total_bits{}; + DecomposerT decomposer{}; + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE + identify_candidates_op_t(key_prefix_t* kth_key_bits, int pass, int total_bits, DecomposerT decomposer) + : kth_key_bits(kth_key_bits) + , pass(pass) + , total_bits(total_bits) + , decomposer(decomposer) + {} + + _CCCL_HOST_DEVICE _CCCL_FORCEINLINE candidate_class operator()(T key) const + { + if (pass <= 0) + { + return candidate_class::candidate; + } + + bit_ordered_type ordered = key; + ordered = RadixSortTwiddle::In(ordered, decomposer); + + // Build the key's prefix using the same funnel shift as set_kth_key_bits + key_prefix_t key_prefix{}; + for (int prefix_pass = 0; prefix_pass < pass; ++prefix_pass) + { + const int start_bit = calc_start_bit(total_bits, prefix_pass); + const int num_bits = + calc_start_bit(total_bits, prefix_pass - 1) - calc_start_bit(total_bits, prefix_pass); + auto extractor = + radix_traits_t::template digit_extractor>(start_bit, num_bits, decomposer); + key_prefix.shift_or(BitsPerPass, static_cast(extractor.Digit(ordered))); + } + + // Compare word-by-word from MSB to LSB + const int total_prefix_bits = pass * BitsPerPass; + const int top_word_idx = (total_prefix_bits - 1) / 32; + const int bits_in_top_word = ((total_prefix_bits - 1) % 32) + 1; + + // Top word may be partially filled + { + unsigned int key_w = key_prefix.words[top_word_idx]; + unsigned int kth_w = kth_key_bits->words[top_word_idx]; + if (bits_in_top_word < 32) + { + const unsigned int mask = (1u << bits_in_top_word) - 1u; + key_w &= mask; + kth_w &= mask; + } + if (key_w < kth_w) + { + return candidate_class::selected; + } + if (key_w > kth_w) + { + return candidate_class::rejected; + } + } + + // Remaining words are fully populated + for (int w = top_word_idx - 1; w >= 0; --w) + { + if (key_prefix.words[w] < kth_key_bits->words[w]) + { + return candidate_class::selected; + } + if (key_prefix.words[w] > kth_key_bits->words[w]) + { + return candidate_class::rejected; + } + } + + return candidate_class::candidate; + } +}; + +template +#if _CCCL_HAS_CONCEPTS() + requires topk_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceTopKKernel( + const KeyInputIteratorT d_keys_in, + const KeyOutputIteratorT d_keys_out, + const ValueInputIteratorT d_values_in, + const ValueOutputIteratorT d_values_out, + KeyInT* const in_buf, + OffsetT* const in_idx_buf, + KeyInT* const out_buf, + OffsetT* const out_idx_buf, + Counter, OffsetT, OutOffsetT>* counter, + OffsetT* const histogram, + const OffsetT num_items, + const OutOffsetT k, + const OffsetT buffer_length, + ExtractBinOpT extract_bin_op, + IdentifyCandidatesOpT identify_candidates_op, + const int pass, + const bool is_last_pass) +{ + static constexpr topk_policy policy = current_policy(); + using agent_topk_policy_t = + agent_topk_policy; + using agent_topk_t = + AgentTopK; + + __shared__ typename agent_topk_t::TempStorage temp_storage; + agent_topk_t( + temp_storage, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + k, + buffer_length, + extract_bin_op, + identify_candidates_op) + .invoke_filter_and_histogram(in_buf, in_idx_buf, out_buf, out_idx_buf, counter, histogram, pass, is_last_pass); +} + +template +#if _CCCL_HAS_CONCEPTS() + requires topk_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceTopKHistogramKernel( + const KeyInputIteratorT d_keys_in, + const KeyOutputIteratorT d_keys_out, + const ValueInputIteratorT d_values_in, + const ValueOutputIteratorT d_values_out, + Counter, OffsetT, OutOffsetT>* counter, + OffsetT* const histogram, + const OffsetT num_items, + const OutOffsetT k, + const OffsetT buffer_length, + ExtractBinOpT extract_bin_op, + const int pass, + const bool is_last_pass) +{ + static constexpr topk_policy policy = current_policy(); + using agent_topk_policy_t = + agent_topk_policy; + using identify_candidates_op_t = NullType; + using agent_topk_t = + AgentTopK; + + __shared__ typename agent_topk_t::TempStorage temp_storage; + agent_topk_t( + temp_storage, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + k, + buffer_length, + extract_bin_op, + identify_candidates_op_t{}) + .invoke_histogram_only(counter, histogram, pass, is_last_pass); +} + +template +#if _CCCL_HAS_CONCEPTS() + requires topk_policy_selector +#endif // _CCCL_HAS_CONCEPTS() +__launch_bounds__(int(current_policy().threads_per_block)) + _CCCL_KERNEL_ATTRIBUTES void DeviceTopKLastFilterKernel( + const KeyInputIteratorT d_keys_in, + const KeyOutputIteratorT d_keys_out, + const ValueInputIteratorT d_values_in, + const ValueOutputIteratorT d_values_out, + KeyInT* const in_buf, + OffsetT* const in_idx_buf, + Counter, OffsetT, OutOffsetT>* counter, + const OffsetT num_items, + const OutOffsetT k, + const OffsetT buffer_length, + IdentifyCandidatesOpT identify_candidates_op, + const int pass) +{ + static constexpr topk_policy policy = current_policy(); + using agent_topk_policy_t = + agent_topk_policy; + using extract_bin_op_t = NullType; + using agent_topk_t = + AgentTopK; + + __shared__ typename agent_topk_t::TempStorage temp_storage; + agent_topk_t( + temp_storage, + d_keys_in, + d_keys_out, + d_values_in, + d_values_out, + num_items, + k, + buffer_length, + extract_bin_op_t{}, + identify_candidates_op) + .invoke_last_filter(in_buf, in_idx_buf, counter, k, pass); +} + +//! @tparam SelectDirection +//! Determines whether to select the smallest or largest K elements. +//! +//! @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 OffsetT +//! Data Type for variables: num_items +//! +//! @tparam OutOffsetT +//! Data Type for variables: k +//! +//! @tparam DecomposerT +//! Implementation detail, do not specify directly, requirements on the content of this type are subject to breaking +//! change. +template