dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

Browse Source 
CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

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

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0
This commit is contained in:
EngineX CI
2026-07-30 09:35:51 +00:00
parent b4d01f481e
commit 56fd68e7dd
8871 changed files with 1454674 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

971
cccl_upstream/cub/cub/block/block_adjacent_difference.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
CUB_NAMESPACE_BEGIN
//! @rst
//! BlockAdjacentDifference provides :ref:`collective <collective-primitives>` 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 <cub/cub.cuh>
//! // or equivalently <cub/block/block_adjacent_difference.cuh>
//!
//! struct CustomDifference
//! {
//! template <typename DataType>
//! __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<int, 128>;
//!
//! // 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 <typename T, int BlockDimX, int BlockDimY = 1, int BlockDimZ = 1>
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 <typename FlagOp, bool HAS_PARAM = BinaryOpHasIdxParam<T, FlagOp>::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 <typename FlagOp>
struct ApplyOp<FlagOp, false>
{
// 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 <int ITEMS_PER_THREAD, typename FlagT, typename FlagOp>
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<FlagOp>::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 <int ITEMS_PER_THREAD, typename FlagT, typename FlagOp>
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<FlagOp>::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 <cub/cub.cuh>
//! // or equivalently <cub/block/block_adjacent_difference.cuh>
//!
//! struct CustomDifference
//! {
//! template <typename DataType>
//! __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<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD, typename OutputType, typename DifferenceOpT>
_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 <cub/cub.cuh>
//! // or equivalently <cub/block/block_adjacent_difference.cuh>
//!
//! struct CustomDifference
//! {
//! template <typename DataType>
//! __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<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD, typename OutputT, typename DifferenceOpT>
_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 <cub/cub.cuh>
//! // or equivalently <cub/block/block_adjacent_difference.cuh>
//!
//! struct CustomDifference
//! {
//! template <typename DataType>
//! __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<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD, typename OutputType, typename DifferenceOpT>
_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 <cub/cub.cuh>
//! // or equivalently <cub/block/block_adjacent_difference.cuh>
//!
//! struct CustomDifference
//! {
//! template <typename DataType>
//! __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<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD, typename OutputType, typename DifferenceOpT>
_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 <cub/cub.cuh>
//! // or equivalently <cub/block/block_adjacent_difference.cuh>
//!
//! struct CustomDifference
//! {
//! template <typename DataType>
//! __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<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD, typename OutputT, typename DifferenceOpT>
_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 <cub/cub.cuh>
//! // or equivalently <cub/block/block_adjacent_difference.cuh>
//!
//! struct CustomDifference
//! {
//! template <typename DataType>
//! __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<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD, typename OutputT, typename DifferenceOpT>
_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 <cub/cub.cuh>
//! // or equivalently <cub/block/block_adjacent_difference.cuh>
//!
//! struct CustomDifference
//! {
//! template <typename DataType>
//! __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<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD, typename OutputT, typename DifferenceOpT>
_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

1234
cccl_upstream/cub/cub/block/block_discontinuity.cuh Normal file
View File

File diff suppressed because it is too large Load Diff

1314
cccl_upstream/cub/cub/block/block_exchange.cuh Normal file
View File

File diff suppressed because it is too large Load Diff

412
cccl_upstream/cub/cub/block/block_histogram.cuh Normal file
View File

@@ -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 [<em>collective</em>](../index.html#sec0) methods for
* constructing block-wide histograms from data samples partitioned across a CUDA thread block.
*/
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/specializations/block_histogram_atomic.cuh>
#include <cub/block/specializations/block_histogram_sort.cuh>
#include <cub/util_ptx.cuh>
#include <cuda/std/__type_traits/conditional.h>
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 <collective-primitives>` methods for
//! constructing block-wide histograms from data samples partitioned across a CUDA thread block.
//!
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! - A `histogram <http://en.wikipedia.org/wiki/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 <cub/cub.cuh> // or equivalently <cub/block/block_histogram.cuh>
//!
//! __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<unsigned char, 128, 4, 256>;
//!
//! // 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 <typename T,
int BlockDimX,
int ItemsPerThread,
int Bins,
BlockHistogramAlgorithm Algorithm = BLOCK_HISTO_SORT,
int BlockDimY = 1,
int BlockDimZ = 1>
class BlockHistogram
{
private:
/// The thread block size in threads
static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ;
/// Internal specialization.
using InternalBlockHistogram =
::cuda::std::_If<Algorithm == BLOCK_HISTO_SORT,
detail::BlockHistogramSort<T, BlockDimX, ItemsPerThread, Bins, BlockDimY, BlockDimZ>,
detail::BlockHistogramAtomic<Bins>>;
/// 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 <cub/cub.cuh> // or equivalently <cub/block/block_histogram.cuh>
//!
//! __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<unsigned char, 128, 4, 256>;
//!
//! // 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 <typename CounterT>
_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 <cub/cub.cuh> // or equivalently <cub/block/block_histogram.cuh>
//!
//! __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<unsigned char, 128, 4, 256>;
//!
//! // 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 <typename CounterT>
_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 <cub/cub.cuh> // or equivalently <cub/block/block_histogram.cuh>
//!
//! __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<unsigned char, 128, 4, 256>;
//!
//! // 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 <typename CounterT>
_CCCL_DEVICE _CCCL_FORCEINLINE void Composite(T (&items)[ItemsPerThread], CounterT histogram[Bins])
{
InternalBlockHistogram(temp_storage).Composite(items, histogram);
}
};
CUB_NAMESPACE_END

1192
cccl_upstream/cub/cub/block/block_load.cuh Normal file
View File

File diff suppressed because it is too large Load Diff

437
cccl_upstream/cub/cub/block/block_load_to_shared.cuh Normal file
View File

@@ -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 <collective-primitives>` method for asynchronously
//! loading data from global to shared memory.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_device.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <thrust/type_traits/is_trivially_relocatable.h>
#include <cuda/__cmath/round_down.h>
#include <cuda/__cmath/round_up.h>
#include <cuda/__memory/address_space.h>
#include <cuda/__memory/align_up.h>
#include <cuda/__memory/is_aligned.h>
#include <cuda/__memory/is_valid_alignment.h>
#include <cuda/__memory/ptr_rebind.h>
#include <cuda/__ptx/instructions/cp_async_bulk.h>
#include <cuda/__ptx/instructions/elect_sync.h>
#include <cuda/__ptx/instructions/mbarrier_arrive.h>
#include <cuda/__ptx/instructions/mbarrier_init.h>
#include <cuda/__ptx/instructions/mbarrier_inval.h>
#include <cuda/__ptx/instructions/mbarrier_wait.h>
#include <cuda/std/__algorithm/max.h>
#include <cuda/std/__algorithm/min.h>
#include <cuda/std/__iterator/data.h>
#include <cuda/std/__iterator/size.h>
#include <cuda/std/cstdint>
#include <cuda/std/span>
#include <nv/target>
CUB_NAMESPACE_BEGIN
namespace detail
{
//! @rst
//! The @c BlockLoadToShared class provides a :ref:`collective <collective-primitives>` 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<T,
//! GmemAlign>(num_items)` and aligned according to `cub::detail::LoadToSharedBufferAlignBytes<T>()`.
//! - 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 <int BlockDimX, int BlockDimY = 1, int BlockDimZ = 1>
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<vec_load_t>(thread_dst) = *::cuda::ptr_rebind<vec_load_t>(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<BlockDimX, BlockDimY, BlockDimZ>&) = delete;
//! @}
_CCCL_DEVICE_API BlockLoadToShared& operator=(const BlockLoadToShared<BlockDimX, BlockDimY, BlockDimZ>&) = 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<T>()` and at least
//! `SharedBufferSizeBytes<T, GmemAlign>(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 <typename T, ::cuda::std::size_t GmemAlign = alignof(T)>
[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE ::cuda::std::span<T>
CopyAsync(::cuda::std::span<char> smem_dst, ::cuda::std::span<const T> gmem_src)
{
static_assert(THRUST_NS_QUALIFIER::is_trivially_relocatable_v<T>);
static_assert(::cuda::__is_valid_alignment<T>(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<int>(sizeof(T)) * static_cast<int>(size(gmem_src));
const auto dst_ptr = data(smem_dst);
const auto src_ptr = ::cuda::ptr_rebind<char>(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<T>()),
"Shared memory needs to be 16 byte aligned.");
_CCCL_ASSERT(
(static_cast<int>(size(smem_dst)) >= cub::detail::LoadToSharedBufferSizeBytes<T, GmemAlign>(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<T>(::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<int>(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<T>(actual_dst_ptr), ::cuda::std::size(gmem_src)};
}
}
// Avoid need to explicitly specify `T` for non-const src.
//! @brief Convenience overload, see `CopyAsync(span<char>, span<const T>)`.
template <typename T, ::cuda::std::size_t GmemAlign = alignof(T)>
[[nodiscard]] _CCCL_DEVICE_API _CCCL_FORCEINLINE ::cuda::std::span<T>
CopyAsync(::cuda::std::span<char> smem_dst, ::cuda::std::span<T> gmem_src)
{
return CopyAsync<T, GmemAlign>(smem_dst, ::cuda::std::span<const T>{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

859
cccl_upstream/cub/cub/block/block_merge_sort.cuh Normal file
View File

@@ -0,0 +1,859 @@
// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/thread/thread_sort.cuh>
#include <cub/util_math.cuh>
#include <cub/util_namespace.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <cuda/__cmath/pow2.h>
#include <cuda/std/__algorithm/min.h>
#include <cuda/std/__type_traits/is_same.h>
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 <typename KeyIt1, typename KeyIt2, typename OffsetT, typename BinaryPred>
_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<OffsetT>(keys1_begin, keys1_end);
// pull copies of the keys before calling binary_pred so proxy references are unwrapped
const detail::it_value_t<KeyIt1> key1 = keys1[mid];
const detail::it_value_t<KeyIt2> key2 = keys2[diag - 1 - mid];
if (binary_pred(key2, key1))
{
keys1_end = mid;
}
else
{
keys1_begin = mid + 1;
}
}
return keys1_begin;
}
namespace detail
{
template <bool Unroll = true, typename KeyIt, typename KeyT, typename CompareOp, int ItemsPerThread>
_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 <bool Unroll = true, typename KeyIt, typename KeyT, typename CompareOp, int ItemsPerThread>
_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<Unroll>(
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 <typename KeyIt, typename KeyT, typename CompareOp, int ItemsPerThread>
_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 <typename KeyIt, typename KeyT, typename CompareOp, int ItemsPerThread>
_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 <cub/cub.cuh> // or equivalently <cub/block/block_merge_sort.cuh>
*
* constexpr int BLOCK_THREADS = 256;
* constexpr int ItemsPerThread = 9;
*
* class BlockMergeSort : public BlockMergeSortStrategy<int,
* cub::NullType,
* BLOCK_THREADS,
* ItemsPerThread,
* BlockMergeSort>
* {
* using BlockMergeSortStrategyT =
* BlockMergeSortStrategy<int,
* cub::NullType,
* BLOCK_THREADS,
* ItemsPerThread,
* BlockMergeSort>;
* 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 <typename KeyT,
typename ValueT,
int NumThreads,
int ItemsPerThread,
typename SynchronizationPolicy,
bool _Unroll = true>
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<ValueT, NullType>;
#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 <typename CompareOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void Sort(KeyT (&keys)[ItemsPerThread], CompareOp compare_op)
{
ValueT items[ItemsPerThread];
Sort<CompareOp, false>(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 <typename CompareOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
Sort(KeyT (&keys)[ItemsPerThread], CompareOp compare_op, int valid_items, KeyT oob_default)
{
ValueT items[ItemsPerThread];
Sort<CompareOp, true>(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 <typename CompareOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
Sort(KeyT (&keys)[ItemsPerThread], ValueT (&items)[ItemsPerThread], CompareOp compare_op)
{
Sort<CompareOp, false>(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 <typename CompareOp, bool IS_LAST_TILE = true>
_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 <typename CompareOp>
_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 <typename CompareOp>
_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 <typename CompareOp>
_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 <typename CompareOp, bool IS_LAST_TILE = true>
_CCCL_DEVICE _CCCL_FORCEINLINE void StableSort(
KeyT (&keys)[ItemsPerThread],
ValueT (&items)[ItemsPerThread],
CompareOp compare_op,
int valid_items,
KeyT oob_default)
{
Sort<CompareOp, IS_LAST_TILE>(keys, items, compare_op, valid_items, oob_default);
}
private:
_CCCL_DEVICE _CCCL_FORCEINLINE void Sync() const
{
static_cast<const SynchronizationPolicy*>(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 <cub/cub.cuh> // or equivalently <cub/block/block_merge_sort.cuh>
*
* struct CustomLess
* {
* template <typename DataType>
* __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<int, 128, 4>;
*
* // 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 <typename KeyT,
int BlockDimX,
int ItemsPerThread,
typename ValueT = NullType,
int BlockDimY = 1,
int BlockDimZ = 1,
bool _Unroll = true>
class BlockMergeSort
: public BlockMergeSortStrategy<
KeyT,
ValueT,
BlockDimX * BlockDimY * BlockDimZ,
ItemsPerThread,
BlockMergeSort<KeyT, BlockDimX, ItemsPerThread, ValueT, BlockDimY, BlockDimZ, _Unroll>,
_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<KeyT, ValueT, BLOCK_THREADS, ItemsPerThread, BlockMergeSort, _Unroll>;
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

1243
cccl_upstream/cub/cub/block/block_radix_rank.cuh Normal file
View File

File diff suppressed because it is too large Load Diff

2191
cccl_upstream/cub/cub/block/block_radix_sort.cuh Normal file
View File

File diff suppressed because it is too large Load Diff

124
cccl_upstream/cub/cub/block/block_raking_layout.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_type.cuh>
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 <typename T, int ThreadsPerBlock>
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

711
cccl_upstream/cub/cub/block/block_reduce.cuh Normal file
View File

@@ -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 <collective-primitives>` methods for
//! computing a parallel reduction of items partitioned across a CUDA thread block.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/specializations/block_reduce_raking.cuh>
#include <cub/block/specializations/block_reduce_raking_commutative_only.cuh>
#include <cub/block/specializations/block_reduce_warp_reductions.cuh>
#include <cub/thread/thread_operators.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <cuda/std/__concepts/same_as.h>
#include <cuda/std/__functional/operations.h>
#include <cuda/std/__fwd/format.h>
#include <cuda/std/__host_stdlib/ostream>
#include <cuda/std/__type_traits/conditional.h>
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 "<unknown BlockReduceAlgorithm>";
}
} // 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<char> CharT>
struct std::formatter<CUB_NS_QUALIFIER::BlockReduceAlgorithm, CharT> : formatter<const CharT*, CharT>
{
template <class FmtCtx>
auto format(const CUB_NS_QUALIFIER::BlockReduceAlgorithm& algo, FmtCtx& ctx) const
{
return formatter<const CharT*, CharT>::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 <collective-primitives>` methods for computing a
//! parallel reduction of items partitioned across a CUDA thread block.
//!
//! Overview
//! +++++++++++++++++++++++++++++++++++++++++++++
//!
//! - A `reduction <http://en.wikipedia.org/wiki/Reduce_(higher-order_function)>`_ (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 <flexible-data-arrangement>` across 128 threads where each thread
//! owns 4 consecutive items.
//!
//! .. code-block:: c++
//!
//! #include <cub/cub.cuh> // or equivalently <cub/block/block_reduce.cuh>
//!
//! __global__ void ExampleKernel(...)
//! {
//! // Specialize BlockReduce for a 1D block of 128 threads of type int
//! using BlockReduce = cub::BlockReduce<int, 128>;
//!
//! // 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 <typename T,
int BlockDimX,
BlockReduceAlgorithm Algorithm = BLOCK_REDUCE_WARP_REDUCTIONS,
int BlockDimY = 1,
int BlockDimZ = 1>
class BlockReduce
{
private:
/// The thread block size in threads
static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ;
using WarpReductions = detail::BlockReduceWarpReductions<T, BlockDimX, BlockDimY, BlockDimZ>;
using WarpReductionsNondeterministic = detail::BlockReduceWarpReductions<T, BlockDimX, BlockDimY, BlockDimZ, false>;
using RakingCommutativeOnly = detail::BlockReduceRakingCommutativeOnly<T, BlockDimX, BlockDimY, BlockDimZ>;
using Raking = detail::BlockReduceRaking<T, BlockDimX, BlockDimY, BlockDimZ>;
/// Internal specialization type
using InternalBlockReduce =
::cuda::std::_If<Algorithm == BLOCK_REDUCE_WARP_REDUCTIONS,
WarpReductions,
::cuda::std::_If<Algorithm == BLOCK_REDUCE_WARP_REDUCTIONS_NONDETERMINISTIC,
WarpReductionsNondeterministic,
::cuda::std::_If<Algorithm == BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY,
RakingCommutativeOnly,
Raking>>>; // 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 <cub/cub.cuh> // or equivalently <cub/block/block_reduce.cuh>
//!
//! __global__ void ExampleKernel(...)
//! {
//! // Specialize BlockReduce for a 1D block of 128 threads of type int
//! using BlockReduce = cub::BlockReduce<int, 128>;
//!
//! // 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 <typename ReductionOp>
_CCCL_DEVICE _CCCL_FORCEINLINE T Reduce(T input, ReductionOp reduction_op)
{
return InternalBlockReduce(temp_storage).template Reduce<true>(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 <flexible-data-arrangement>` across 128 threads where each thread owns
//! 4 consecutive items.
//!
//! .. code-block:: c++
//!
//! #include <cub/cub.cuh> // or equivalently <cub/block/block_reduce.cuh>
//!
//! __global__ void ExampleKernel(...)
//! {
//! // Specialize BlockReduce for a 1D block of 128 threads of type int
//! using BlockReduce = cub::BlockReduce<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD, typename ReductionOp>
_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 thread<sub>0</sub>.
//! - @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 <cub/cub.cuh> // or equivalently <cub/block/block_reduce.cuh>
//!
//! __global__ void ExampleKernel(int num_valid, ...)
//! {
//! // Specialize BlockReduce for a 1D block of 128 threads of type int
//! using BlockReduce = cub::BlockReduce<int, 128>;
//!
//! // 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 <typename ReductionOp>
_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<true>(input, num_valid, reduction_op);
}
else
{
return InternalBlockReduce(temp_storage).template Reduce<false>(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 <cub/cub.cuh> // or equivalently <cub/block/block_reduce.cuh>
//!
//! __global__ void ExampleKernel(...)
//! {
//! // Specialize BlockReduce for a 1D block of 128 threads of type int
//! using BlockReduce = cub::BlockReduce<int, 128>;
//!
//! // 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<true>(input, BLOCK_THREADS);
}
//! @rst
//! Computes a block-wide reduction for thread<sub>0</sub> 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 <flexible-data-arrangement>` across 128 threads where each thread owns
//! 4 consecutive items.
//!
//! .. code-block:: c++
//!
//! #include <cub/cub.cuh> // or equivalently <cub/block/block_reduce.cuh>
//!
//! __global__ void ExampleKernel(...)
//! {
//! // Specialize BlockReduce for a 1D block of 128 threads of type int
//! using BlockReduce = cub::BlockReduce<int, 128>;
//!
//! // 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 <int ITEMS_PER_THREAD>
_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 <cub/cub.cuh> // or equivalently <cub/block/block_reduce.cuh>
//!
//! __global__ void ExampleKernel(int num_valid, ...)
//! {
//! // Specialize BlockReduce for a 1D block of 128 threads of type int
//! using BlockReduce = cub::BlockReduce<int, 128>;
//!
//! // 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<true>(input, num_valid);
}
else
{
return InternalBlockReduce(temp_storage).template Sum<false>(input, num_valid);
}
}
//! @}
};
CUB_NAMESPACE_END

436
cccl_upstream/cub/cub/block/block_run_length_decode.cuh Normal file
View File

@@ -0,0 +1,436 @@
// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.
// SPDX-License-Identifier: BSD-3
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/block_scan.cuh>
#include <cub/thread/thread_search.cuh>
#include <cub/util_math.cuh>
#include <cub/util_namespace.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <cuda/std/__algorithm/min.h>
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<RunItemT, BlockDimX, RunsPerThread, DecodedItemsPerThread>;
//!
//! // 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 <typename ItemT,
int BlockDimX,
int RunsPerThread,
int DecodedItemsPerThread,
typename DecodedOffsetT = uint32_t,
int BlockDimY = 1,
int BlockDimZ = 1>
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<DecodedOffsetT, BlockDimX, BLOCK_SCAN_RAKING_MEMOIZE, BlockDimY, BlockDimZ>;
/// 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 <typename RunLengthT, typename TotalDecodedSizeT>
_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 <typename UserRunOffsetT>
_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 <typename RunLengthT, typename TotalDecodedSizeT>
_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 <typename UserRunOffsetT>
_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 <int MAX_NUM_ITEMS, typename InputIteratorT, typename OffsetT, typename T>
_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<MAX_NUM_ITEMS>::VALUE; i++)
{
OffsetT mid = cub::MidPoint<OffsetT>(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 <typename RunOffsetT>
_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<RunOffsetT>(linear_tid) * static_cast<RunOffsetT>(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 <typename RunLengthT, typename TotalDecodedSizeT>
_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<DecodedOffsetT>(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<TotalDecodedSizeT>(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 <typename RelativeOffsetT>
_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 <thread_decoded_offset> is already beyond the total decoded size, it will be assigned to the
// last run
RunOffsetT assigned_run =
StaticUpperBound<BLOCK_RUNS>(temp_storage.runs.run_offsets, BLOCK_RUNS, thread_decoded_offset)
- static_cast<RunOffsetT>(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

2277
cccl_upstream/cub/cub/block/block_scan.cuh Normal file
View File

File diff suppressed because it is too large Load Diff

340
cccl_upstream/cub/cub/block/block_shuffle.cuh Normal file
View File

@@ -0,0 +1,340 @@
// 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::BlockShuffle class provides :ref:`collective <collective-primitives>` methods for shuffling data
//! partitioned across a CUDA thread block.
#pragma once
#include <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
CUB_NAMESPACE_BEGIN
//! @rst
//! The BlockShuffle class provides :ref:`collective <collective-primitives>`
//! methods for shuffling data partitioned across a CUDA thread block.
//!
//! Overview
//! ++++++++++++++++
//!
//! It is commonplace for blocks of threads to rearrange data items between threads.
//! The BlockShuffle abstraction allows threads to efficiently shift items either
//! (a) up to their successor or
//! (b) down to their predecessor
//!
//! @endrst
//!
//! @tparam T
//! The data type to be exchanged.
//!
//! @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 <typename T, int BlockDimX, int BlockDimY = 1, int BlockDimZ = 1>
class BlockShuffle
{
private:
static constexpr int BLOCK_THREADS = BlockDimX * BlockDimY * BlockDimZ;
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;
/// Shared memory storage layout type (last element from each thread's input)
using _TempStorage = T[BLOCK_THREADS];
public:
/// \smemstorage{BlockShuffle}
struct TempStorage : Uninitialized<_TempStorage>
{};
private:
/// 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:
//! @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 BlockShuffle()
: 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 BlockShuffle(TempStorage& temp_storage)
: temp_storage(temp_storage.Alias())
, linear_tid(RowMajorTid(BlockDimX, BlockDimY, BlockDimZ))
{}
//! @}
//! @name Shuffle movement
//! @{
//! @rst
//! Each *thread*\ :sub:`i` obtains the ``input`` provided by *thread*\ :sub:`i + distance`.
//! The offset ``distance`` may be negative.
//!
//! .. versionadded:: 2.2.0
//! First appears in CUDA Toolkit 12.3.
//!
//! - @smemreuse
//! @endrst
//!
//! @param[in] input
//! @rst
//! The input item from the calling thread (*thread*\ :sub:`i`)
//! @endrst
//!
//! @param[out] output
//! @rst
//! The ``input`` item from the successor (or predecessor) thread
//! *thread*\ :sub:`i + distance` (may be aliased to ``input``).
//! This value is only updated for for *thread*\ :sub:`i` when
//! ``0 <= (i + distance) < BLOCK_THREADS - 1``
//! @endrst
//!
//! @param[in] distance
//! Offset distance (may be negative)
_CCCL_DEVICE _CCCL_FORCEINLINE void Offset(T input, T& output, int distance = 1)
{
temp_storage[linear_tid] = input;
__syncthreads();
const int offset_tid = static_cast<int>(linear_tid) + distance;
if ((offset_tid >= 0) && (offset_tid < BLOCK_THREADS))
{
output = temp_storage[static_cast<size_t>(offset_tid)];
}
}
//! @rst
//! Each *thread*\ :sub:`i` obtains the ``input`` provided by *thread*\ :sub:`i + distance`.
//!
//! .. versionadded:: 2.2.0
//! First appears in CUDA Toolkit 12.3.
//!
//! - @smemreuse
//! @endrst
//!
//! @param[in] input
//! The calling thread's input item
//!
//! @param[out] output
//! @rst
//! The ``input`` item from thread
//! *thread*\ :sub:`(i + distance>) % BLOCK_THREADS` (may be aliased to ``input``).
//! This value is not updated for *thread*\ :sub:`BLOCK_THREADS - 1`.
//! @endrst
//!
//! @param[in] distance
//! Offset distance (`0 < distance < `BLOCK_THREADS`)
_CCCL_DEVICE _CCCL_FORCEINLINE void Rotate(T input, T& output, unsigned int distance = 1)
{
temp_storage[linear_tid] = input;
__syncthreads();
unsigned int offset = linear_tid + distance;
if (offset >= BLOCK_THREADS)
{
offset -= BLOCK_THREADS;
}
output = temp_storage[offset];
}
//! @rst
//! The thread block rotates its :ref:`blocked arrangement <flexible-data-arrangement>` of
//! ``input`` items, shifting it up by one item.
//!
//! .. versionadded:: 2.2.0
//! First appears in CUDA Toolkit 12.3.
//!
//! - @blocked
//! - @granularity
//! - @smemreuse
//! @endrst
//!
//! @param[in] input
//! The calling thread's input items
//!
//! @param[out] prev
//! @rst
//! The corresponding predecessor items (may be aliased to ``input``).
//! The item ``prev[0]`` is not updated for *thread*\ :sub:`0`.
//! @endrst
template <int ITEMS_PER_THREAD>
_CCCL_DEVICE _CCCL_FORCEINLINE void Up(T (&input)[ITEMS_PER_THREAD], T (&prev)[ITEMS_PER_THREAD])
{
temp_storage[linear_tid] = input[ITEMS_PER_THREAD - 1];
__syncthreads();
_CCCL_PRAGMA_UNROLL_FULL()
for (int ITEM = ITEMS_PER_THREAD - 1; ITEM > 0; --ITEM)
{
prev[ITEM] = input[ITEM - 1];
}
if (linear_tid > 0)
{
prev[0] = temp_storage[linear_tid - 1];
}
}
//! @rst
//! The thread block rotates its :ref:`blocked arrangement <flexible-data-arrangement>`
//! of ``input`` items, shifting it up by one item. All threads receive the ``input`` provided by
//! *thread*\ :sub:`BLOCK_THREADS - 1`.
//!
//! .. versionadded:: 2.2.0
//! First appears in CUDA Toolkit 12.3.
//!
//! - @blocked
//! - @granularity
//! - @smemreuse
//! @endrst
//!
//! @param[in] input
//! The calling thread's input items
//!
//! @param[out] prev
//! @rst
//! The corresponding predecessor items (may be aliased to ``input``).
//! The item ``prev[0]`` is not updated for *thread*\ :sub:`0`.
//! @endrst
//!
//! @param[out] block_suffix
//! @rst
//! The item ``input[ITEMS_PER_THREAD - 1]`` from *thread*\ :sub:`BLOCK_THREADS - 1`, provided to all threads
//! @endrst
template <int ITEMS_PER_THREAD>
_CCCL_DEVICE _CCCL_FORCEINLINE void Up(T (&input)[ITEMS_PER_THREAD], T (&prev)[ITEMS_PER_THREAD], T& block_suffix)
{
Up(input, prev);
block_suffix = temp_storage[BLOCK_THREADS - 1];
}
//! @rst
//! The thread block rotates its :ref:`blocked arrangement <flexible-data-arrangement>`
//! of ``input`` items, shifting it down by one item.
//!
//! .. versionadded:: 2.2.0
//! First appears in CUDA Toolkit 12.3.
//!
//! - @blocked
//! - @granularity
//! - @smemreuse
//! @endrst
//!
//! @param[in] input
//! The calling thread's input items
//!
//! @param[out] prev
//! @rst
//! The corresponding predecessor items (may be aliased to ``input``).
//! The value ``prev[0]`` is not updated for *thread*\ :sub:`BLOCK_THREADS - 1`.
//! @endrst
template <int ITEMS_PER_THREAD>
_CCCL_DEVICE _CCCL_FORCEINLINE void Down(T (&input)[ITEMS_PER_THREAD], T (&prev)[ITEMS_PER_THREAD])
{
temp_storage[linear_tid] = input[0];
__syncthreads();
_CCCL_PRAGMA_UNROLL_FULL()
for (int ITEM = 0; ITEM < ITEMS_PER_THREAD - 1; ITEM++)
{
prev[ITEM] = input[ITEM + 1];
}
if (linear_tid < BLOCK_THREADS - 1)
{
prev[ITEMS_PER_THREAD - 1] = temp_storage[linear_tid + 1];
}
}
//! @rst
//! The thread block rotates its :ref:`blocked arrangement <flexible-data-arrangement>` of input items,
//! shifting it down by one item. All threads receive ``input[0]`` provided by *thread*\ :sub:`0`.
//!
//! .. versionadded:: 2.2.0
//! First appears in CUDA Toolkit 12.3.
//!
//! - @blocked
//! - @granularity
//! - @smemreuse
//! @endrst
//!
//! @param[in] input
//! The calling thread's input items
//!
//! @param[out] prev
//! @rst
//! The corresponding predecessor items (may be aliased to ``input``).
//! The value ``prev[0]`` is not updated for *thread*\ :sub:`BLOCK_THREADS - 1`.
//! @endrst
//!
//! @param[out] block_prefix
//! @rst
//! The item ``input[0]`` from *thread*\ :sub:`0`, provided to all threads
//! @endrst
template <int ITEMS_PER_THREAD>
_CCCL_DEVICE _CCCL_FORCEINLINE void Down(T (&input)[ITEMS_PER_THREAD], T (&prev)[ITEMS_PER_THREAD], T& block_prefix)
{
Down(input, prev);
block_prefix = temp_storage[0];
}
//! @}
};
CUB_NAMESPACE_END

960
cccl_upstream/cub/cub/block/block_store.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/block_exchange.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <cuda/std/__concepts/same_as.h>
#include <cuda/std/__fwd/format.h>
#include <cuda/std/__host_stdlib/ostream>
#include <cuda/std/__memory/is_sufficiently_aligned.h>
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 <typename T, int ItemsPerThread, typename OutputIteratorT>
_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 <typename T, int ItemsPerThread, typename OutputIteratorT>
_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 <typename T, int ItemsPerThread>
_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<T, VEC_SIZE>::Type;
// Add the alignment check to ensure the vectorized storing can proceed.
if (::cuda::std::is_sufficiently_aligned<alignof(Vector)>(block_ptr))
{
// Alias global pointer
Vector* block_ptr_vectors = reinterpret_cast<Vector*>(const_cast<T*>(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<T*>(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 <int BLOCK_THREADS, typename T, int ItemsPerThread, typename OutputIteratorT>
_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 <int BLOCK_THREADS, typename T, int ItemsPerThread, typename OutputIteratorT>
_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 <typename T, int ItemsPerThread, typename OutputIteratorT>
_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 <typename T, int ItemsPerThread, typename OutputIteratorT>
_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 <flexible-data-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 <flexible-data-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 <flexible-data-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 <flexible-data-arrangement>` is locally
//! transposed and then efficiently written to memory as a :ref:`striped arrangement <flexible-data-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 <flexible-data-arrangement>` is locally
//! transposed and then efficiently written to memory as a
//! :ref:`warp-striped arrangement <flexible-data-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 <flexible-data-arrangement>` is locally
//! transposed and then efficiently written to memory as a
//! :ref:`warp-striped arrangement <flexible-data-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 "<unknown BlockStoreAlgorithm>";
}
} // 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<char> CharT>
struct std::formatter<CUB_NS_QUALIFIER::BlockStoreAlgorithm, CharT> : formatter<const CharT*, CharT>
{
template <class FmtCtx>
auto format(const CUB_NS_QUALIFIER::BlockStoreAlgorithm& algo, FmtCtx& ctx) const
{
return formatter<const CharT*, CharT>::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 <collective-primitives>` data movement
//! methods for writing a :ref:`blocked arrangement <flexible-data-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 <flexible-data-arrangement>` of data is written directly to memory.
//! #. :cpp:enumerator:`cub::BLOCK_STORE_STRIPED`:
//! A :ref:`striped arrangement <flexible-data-arrangement>` of data is written directly to memory.
//! #. :cpp:enumerator:`cub::BLOCK_STORE_VECTORIZE`:
//! A :ref:`blocked arrangement <flexible-data-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 <flexible-data-arrangement>` is locally transposed into
//! a :ref:`striped arrangement <flexible-data-arrangement>` which is then written to memory.
//! #. :cpp:enumerator:`cub::BLOCK_STORE_WARP_TRANSPOSE`:
//! A :ref:`blocked arrangement <flexible-data-arrangement>` is locally transposed into
//! a :ref:`warp-striped arrangement <flexible-data-arrangement>` which is then written to memory.
//! #. :cpp:enumerator:`cub::BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED`:
//! A :ref:`blocked arrangement <flexible-data-arrangement>` is locally transposed into
//! a :ref:`warp-striped arrangement <flexible-data-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 <cub/cub.cuh> // or equivalently <cub/block/block_store.cuh>
//!
//! __global__ void ExampleKernel(int *d_data, ...)
//! {
//! // Specialize BlockStore for a 1D block of 128 threads owning 4 integer items each
//! using BlockStore = cub::BlockStore<int, 128, 4, BLOCK_STORE_WARP_TRANSPOSE>;
//!
//! // 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 <typename T,
int BlockDimX,
int ItemsPerThread,
BlockStoreAlgorithm Algorithm = BLOCK_STORE_DIRECT,
int BlockDimY = 1,
int BlockDimZ = 1>
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<T, BlockDimX, ItemsPerThread, Algorithm == BLOCK_STORE_WARP_TRANSPOSE_TIMESLICED, BlockDimY, BlockDimZ>;
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 <cub/cub.cuh> // or equivalently <cub/block/block_store.cuh>
//!
//! __global__ void ExampleKernel(int *d_data, ...)
//! {
//! // Specialize BlockStore for a 1D block of 128 threads owning 4 integer items each
//! using BlockStore = cub::BlockStore<int, 128, 4, BLOCK_STORE_WARP_TRANSPOSE>;
//!
//! // 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 <typename OutputIteratorT>
_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<BLOCK_THREADS>(linear_tid, block_itr, items);
}
else if constexpr (Algorithm == BLOCK_STORE_VECTORIZE)
{
if constexpr (::cuda::std::contiguous_iterator<OutputIteratorT> && ::cuda::std::__can_to_address<OutputIteratorT>)
{
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<BLOCK_THREADS>(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 <cub/cub.cuh> // or equivalently <cub/block/block_store.cuh>
//!
//! __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<int, 128, 4, BLOCK_STORE_WARP_TRANSPOSE>;
//!
//! // 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 <typename OutputIteratorT>
_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<BLOCK_THREADS>(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<BLOCK_THREADS>(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 <class Policy, class It, class T = cub::detail::it_value_t<It>>
struct BlockStoreType
{
using type = cub::BlockStore<T, Policy::BLOCK_THREADS, Policy::ITEMS_PER_THREAD, Policy::STORE_ALGORITHM>;
};
#endif // _CCCL_DOXYGEN_INVOKED
CUB_NAMESPACE_END

89
cccl_upstream/cub/cub/block/block_topk.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/specializations/block_topk_air.cuh>
#include <cub/device/dispatch/dispatch_common.cuh>
#include <cub/util_type.cuh>
CUB_NAMESPACE_BEGIN
namespace detail
{
// TODO (elstehle): Add documentation
template <typename KeyT, int BlockDimX, int ItemsPerThread, typename ValueT = NullType>
class block_topk
{
private:
using internal_block_topk_t = block_topk_air<KeyT, BlockDimX, ItemsPerThread, ValueT>;
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 <bool IsFullTile>
_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<detail::topk::select::max, IsFullTile>(keys, values, k, num_valid, begin_bit, end_bit);
}
template <bool IsFullTile>
_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<detail::topk::select::max, IsFullTile>(keys, k, num_valid, begin_bit, end_bit);
}
template <bool IsFullTile>
_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<detail::topk::select::min, IsFullTile>(keys, values, k, num_valid, begin_bit, end_bit);
}
template <bool IsFullTile>
_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<detail::topk::select::min, IsFullTile>(keys, k, num_valid, begin_bit, end_bit);
}
};
} // namespace detail
CUB_NAMESPACE_END

580
cccl_upstream/cub/cub/block/radix_rank_sort_operations.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/type_traits.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <thrust/type_traits/integer_sequence.h>
#include <cuda/__bit/bitfield.h>
#include <cuda/__type_traits/is_floating_point.h>
#include <cuda/__utility/static_for.h>
#include <cuda/std/__algorithm/min.h>
#include <cuda/std/__functional/invoke.h>
#include <cuda/std/__type_traits/enable_if.h>
#include <cuda/std/__type_traits/integral_constant.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__type_traits/remove_cv.h>
#include <cuda/std/__type_traits/void_t.h>
#include <cuda/std/cstdint>
#include <cuda/std/tuple>
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 <typename KeyT, bool IsFP = ::cuda::is_floating_point_v<KeyT>>
struct BaseDigitExtractor
{
using TraitsT = Traits<KeyT>;
using UnsignedBits = typename TraitsT::UnsignedBits;
static _CCCL_HOST_DEVICE _CCCL_FORCEINLINE UnsignedBits ProcessFloatMinusZero(UnsignedBits key)
{
return key;
}
};
template <typename KeyT>
struct BaseDigitExtractor<KeyT, true>
{
using TraitsT = Traits<KeyT>;
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 <typename KeyT>
struct BFEDigitExtractor : BaseDigitExtractor<KeyT>
{
using typename BaseDigitExtractor<KeyT>::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 <typename KeyT>
struct ShiftDigitExtractor : BaseDigitExtractor<KeyT>
{
using typename BaseDigitExtractor<KeyT>::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 <class T>
_CCCL_HOST_DEVICE T& operator()(T& key) const
{
return key;
}
};
template <class F, class... Ts, ::cuda::std::size_t... Is>
_CCCL_HOST_DEVICE void
for_each_member_impl(F f, const ::cuda::std::tuple<Ts&...>& tpl, ::cuda::std::index_sequence<Is...>)
{
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<Is>(tpl)), dummy) = ... = 0);
}
template <class F, class DecomposerT, class T>
_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<decltype(tuple_of_refs)>>;
for_each_member_impl(f, tuple_of_refs, ::cuda::std::make_index_sequence<tuple_size>{});
}
namespace radix
{
// True for types that can be converted to bit ordered values using cub::Traits<T>::UnsignedBits (and TwiddleIn/Out)
template <class T, class = void>
inline constexpr bool can_twiddle = false;
template <class T>
inline constexpr bool can_twiddle<T, ::cuda::std::void_t<typename Traits<T>::UnsignedBits>> = true;
template <class T>
inline constexpr bool can_twiddle_tuple_refs = false;
template <class... Ts>
inline constexpr bool can_twiddle_tuple_refs<::cuda::std::tuple<Ts&...>> = (can_twiddle<Ts> && ...);
template <class KeyT, class DecomposerT>
inline constexpr bool decomposer_check = can_twiddle_tuple_refs<::cuda::std::invoke_result_t<DecomposerT, KeyT&>>;
// SFINAE-friendly version of decomposer_check_t: true iff DecomposerT is callable
// with KeyT& and returns a tuple of references to fundamental types.
template <class KeyT, class DecomposerT, class = void>
inline constexpr bool is_valid_decomposer = false;
template <class KeyT, class DecomposerT>
inline constexpr bool
is_valid_decomposer<KeyT, DecomposerT, ::cuda::std::void_t<::cuda::std::invoke_result_t<DecomposerT, KeyT&>>> =
can_twiddle_tuple_refs<::cuda::std::invoke_result_t<DecomposerT, KeyT&>>;
template <class T>
struct bit_ordered_conversion_policy_t
{
using bit_ordered_type = typename Traits<T>::UnsignedBits;
static _CCCL_HOST_DEVICE bit_ordered_type to_bit_ordered(detail::identity_decomposer_t, bit_ordered_type val)
{
return Traits<T>::TwiddleIn(val);
}
static _CCCL_HOST_DEVICE bit_ordered_type from_bit_ordered(detail::identity_decomposer_t, bit_ordered_type val)
{
return Traits<T>::TwiddleOut(val);
}
};
template <class T>
struct bit_ordered_inversion_policy_t
{
using bit_ordered_type = typename Traits<T>::UnsignedBits;
static _CCCL_HOST_DEVICE bit_ordered_type inverse(detail::identity_decomposer_t, bit_ordered_type val)
{
return ~val;
}
};
template <class T, bool = can_twiddle<T>>
struct traits_t
{
using bit_ordered_type = typename Traits<T>::UnsignedBits;
using bit_ordered_conversion_policy = bit_ordered_conversion_policy_t<T>;
using bit_ordered_inversion_policy = bit_ordered_inversion_policy_t<T>;
template <class FundamentalExtractorT, class /* DecomposerT */>
using digit_extractor_t = FundamentalExtractorT;
static _CCCL_HOST_DEVICE bit_ordered_type min_raw_binary_key(detail::identity_decomposer_t)
{
return Traits<T>::LOWEST_KEY;
}
static _CCCL_HOST_DEVICE bit_ordered_type max_raw_binary_key(detail::identity_decomposer_t)
{
return Traits<T>::MAX_KEY;
}
static _CCCL_HOST_DEVICE int default_end_bit(detail::identity_decomposer_t)
{
return sizeof(T) * 8;
}
template <class FundamentalExtractorT>
static _CCCL_HOST_DEVICE digit_extractor_t<FundamentalExtractorT, detail::identity_decomposer_t>
digit_extractor(int begin_bit, int num_bits, detail::identity_decomposer_t)
{
return FundamentalExtractorT(begin_bit, num_bits);
}
};
template <class T, bool CanTwiddle>
struct traits_t<T&, CanTwiddle> : traits_t<T, CanTwiddle>
{};
template <class DecomposerT>
struct min_raw_binary_key_f
{
DecomposerT decomposer;
template <class T>
_CCCL_HOST_DEVICE void operator()(T& field)
{
using traits = traits_t<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<bit_ordered_type&>(field) = traits::min_raw_binary_key(detail::identity_decomposer_t{});
}
};
template <class DecomposerT, class T>
_CCCL_HOST_DEVICE void min_raw_binary_key(DecomposerT decomposer, T& aggregate)
{
detail::for_each_member(min_raw_binary_key_f<DecomposerT>{decomposer}, decomposer, aggregate);
}
template <class DecomposerT>
struct max_raw_binary_key_f
{
DecomposerT decomposer;
template <class T>
_CCCL_HOST_DEVICE void operator()(T& field)
{
using traits = traits_t<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<bit_ordered_type&>(field) = traits::max_raw_binary_key(detail::identity_decomposer_t{});
}
};
template <class DecomposerT, class T>
_CCCL_HOST_DEVICE void max_raw_binary_key(DecomposerT decomposer, T& aggregate)
{
detail::for_each_member(max_raw_binary_key_f<DecomposerT>{decomposer}, decomposer, aggregate);
}
template <class DecomposerT>
struct to_bit_ordered_f
{
DecomposerT decomposer;
template <class T>
_CCCL_HOST_DEVICE void operator()(T& field)
{
using traits = traits_t<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<bit_ordered_type&>(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 <class DecomposerT, class T>
_CCCL_HOST_DEVICE void to_bit_ordered(DecomposerT decomposer, T& aggregate)
{
detail::for_each_member(to_bit_ordered_f<DecomposerT>{decomposer}, decomposer, aggregate);
}
template <class DecomposerT>
struct from_bit_ordered_f
{
DecomposerT decomposer;
template <class T>
_CCCL_HOST_DEVICE void operator()(T& field)
{
using traits = traits_t<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<bit_ordered_type&>(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 <class DecomposerT, class T>
_CCCL_HOST_DEVICE void from_bit_ordered(DecomposerT decomposer, T& aggregate)
{
detail::for_each_member(from_bit_ordered_f<DecomposerT>{decomposer}, decomposer, aggregate);
}
struct inverse_f
{
template <class T>
_CCCL_HOST_DEVICE void operator()(T& field)
{
using traits = traits_t<T>;
using bit_ordered_type = typename traits::bit_ordered_type;
auto& ordered_field = reinterpret_cast<bit_ordered_type&>(field);
ordered_field = ~ordered_field;
}
};
template <class DecomposerT, class T>
_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 <class T>
_CCCL_HOST_DEVICE void operator()(T& /* field */)
{
result += sizeof(T) * 8;
}
};
template <class DecomposerT, class T>
_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 <class T>
_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<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<T>::ProcessFloatMinusZero(reinterpret_cast<bit_ordered_type&>(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 <class DecomposerT, class T>
_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 <class DecomposerT>
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 <class T>
_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 <class DecomposerT, class T>
static _CCCL_HOST_DEVICE T to_bit_ordered(DecomposerT decomposer, T val)
{
detail::radix::to_bit_ordered(decomposer, val);
return val;
}
template <class DecomposerT, class T>
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 <class DecomposerT, class T>
static _CCCL_HOST_DEVICE T inverse(DecomposerT decomposer, T val)
{
detail::radix::inverse(decomposer, val);
return val;
}
};
template <class T>
struct traits_t<T, false /* is_fundamental */>
{
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 <class FundamentalExtractorT, class DecomposerT>
using digit_extractor_t = custom_digit_extractor_t<DecomposerT>;
template <class DecomposerT>
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 <class DecomposerT>
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 <class DecomposerT>
static _CCCL_HOST_DEVICE int default_end_bit(DecomposerT decomposer)
{
T aggregate{};
return detail::radix::default_end_bit(decomposer, aggregate);
}
template <class FundamentalExtractorT, class DecomposerT>
static _CCCL_HOST_DEVICE digit_extractor_t<FundamentalExtractorT, DecomposerT>
digit_extractor(int begin_bit, int num_bits, DecomposerT decomposer)
{
return custom_digit_extractor_t<DecomposerT>(decomposer, begin_bit, num_bits);
}
};
} // namespace radix
} // namespace detail
#endif // _CCCL_DOXYGEN_INVOKED
//! Twiddling keys for radix sort
template <bool IS_DESCENDING, typename KeyT>
struct RadixSortTwiddle
{
private:
using traits = detail::radix::traits_t<KeyT>;
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 <class DecomposerT = detail::identity_decomposer_t>
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 <class DecomposerT = detail::identity_decomposer_t>
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 <class DecomposerT = detail::identity_decomposer_t>
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

62
cccl_upstream/cub/cub/block/specializations/block_histogram_atomic.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
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 <int Bins>
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 <typename T, typename CounterT, int ITEMS_PER_THREAD>
_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

209
cccl_upstream/cub/cub/block/specializations/block_histogram_sort.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/block_discontinuity.cuh>
#include <cub/block/block_radix_sort.cuh>
#include <cub/util_ptx.cuh>
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 <typename T, int BlockDimX, int ItemsPerThread, int Bins, int BlockDimY, int BlockDimZ>
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<T,
BlockDimX,
ItemsPerThread,
NullType,
4,
true,
BLOCK_SCAN_WARP_SCANS,
cudaSharedMemBankSizeFourByte,
BlockDimY,
BlockDimZ>;
// Parameterize BlockDiscontinuity type for our thread block
using BlockDiscontinuityT = BlockDiscontinuity<T, BlockDimX, BlockDimY, BlockDimZ>;
/// 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 <typename CounterT>
_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

230
cccl_upstream/cub/cub/block/specializations/block_reduce_raking.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/block_raking_layout.cuh>
#include <cub/thread/thread_reduce.cuh>
#include <cub/util_ptx.cuh>
#include <cub/warp/warp_reduce.cuh>
#include <cuda/__cmath/pow2.h>
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 <typename T, int BlockDimX, int BlockDimY, int BlockDimZ>
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<T, BLOCK_THREADS>;
/// WarpReduce utility type
using WarpReduce = typename WarpReduce<T, BlockRakingLayout::RAKING_THREADS>::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
* <b>[<em>lane</em><sub>0</sub> only]</b> Warp-wide aggregate reduction of input items
*
* @param[in] num_valid
* Number of valid elements (may be less than BLOCK_THREADS)
*/
template <bool IS_FULL_TILE, typename ReductionOp, int ITERATION>
_CCCL_DEVICE _CCCL_FORCEINLINE T RakingReduction(
ReductionOp reduction_op, T* raking_segment, T partial, int num_valid, constant_t<ITERATION> /*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<IS_FULL_TILE>(reduction_op, raking_segment, partial, num_valid, constant_t<ITERATION + 1>());
}
/**
* @param[in] reduction_op
* Binary reduction operator
*
* @param[in] partial
* <b>[<em>lane</em><sub>0</sub> only]</b> Warp-wide aggregate reduction of input items
*
* @param[in] num_valid
* Number of valid elements (may be less than BLOCK_THREADS)
*/
template <bool IS_FULL_TILE, typename ReductionOp>
_CCCL_DEVICE _CCCL_FORCEINLINE T RakingReduction(
ReductionOp /*reduction_op*/,
T* /*raking_segment*/,
T partial,
int /*num_valid*/,
constant_t<SEGMENT_LENGTH> /*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 thread<sub>0</sub>.
*
* @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 <bool IS_FULL_TILE, typename ReductionOp>
_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<IS_FULL_TILE>(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<IS_FULL_TILE>(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<unsigned int>((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 thread<sub>0</sub>.
*
* @param[in] partial
* Calling thread's input partial reductions
*
* @param[in] num_valid
* Number of valid elements (may be less than BLOCK_THREADS)
*/
template <bool IS_FULL_TILE>
_CCCL_DEVICE _CCCL_FORCEINLINE T Sum(T partial, int num_valid)
{
::cuda::std::plus<> reduction_op;
return Reduce<IS_FULL_TILE>(partial, num_valid, reduction_op);
}
};
} // namespace detail
CUB_NAMESPACE_END

207
cccl_upstream/cub/cub/block/specializations/block_reduce_raking_commutative_only.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/specializations/block_reduce_raking.cuh>
#include <cub/thread/thread_reduce.cuh>
#include <cub/util_ptx.cuh>
#include <cub/warp/warp_reduce.cuh>
#include <cuda/std/span>
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 <typename T, int BlockDimX, int BlockDimY, int BlockDimZ>
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<T, BlockDimX, BlockDimY, BlockDimZ>;
/// 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<T, RAKING_THREADS>;
/// Layout type for padded thread block raking grid
using BlockRakingLayout = BlockRakingLayout<T, SHARING_THREADS>;
/// 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 thread<sub>0</sub>.
*
* @param[in] partial
* Calling thread's input partial reductions
*
* @param[in] num_valid
* Number of valid elements (may be less than BLOCK_THREADS)
*/
template <bool FULL_TILE>
_CCCL_DEVICE _CCCL_FORCEINLINE T Sum(T partial, int num_valid)
{
if (USE_FALLBACK || !FULL_TILE)
{
return FallBack(temp_storage.fallback_storage).template Sum<FULL_TILE>(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<T, SEGMENT_LENGTH>(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 thread<sub>0</sub>.
*
* @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 <bool FULL_TILE, typename ReductionOp>
_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<FULL_TILE>(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<T, SEGMENT_LENGTH>(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

260
cccl_upstream/cub/cub/block/specializations/block_reduce_warp_reductions.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/uninitialized_copy.cuh>
#include <cub/util_ptx.cuh>
#include <cub/warp/warp_reduce.cuh>
#include <cuda/__cmath/ceil_div.h>
#include <cuda/__ptx/instructions/get_sreg.h>
#include <cuda/atomic>
#include <cuda/std/__algorithm/min.h>
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 <typename T, int BlockDimX, int BlockDimY, int BlockDimZ, bool IsDeterministic = true>
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<T, logical_warp_size>::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<int>(::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 <typename ReductionOp>
_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<T, ::cuda::thread_scope_block> 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 <bool FullTile, typename ReductionOp>
_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 <bool FullTile>
_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<FullTile>(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 <bool FullTile, typename ReductionOp>
_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<FullTile>(reduction_op, warp_aggregate, num_valid);
}
else
{
return ApplyWarpAggregatesNonDeterministic(reduction_op, warp_aggregate);
}
}
};
} // namespace detail
CUB_NAMESPACE_END

766
cccl_upstream/cub/cub/block/specializations/block_scan_raking.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/block_raking_layout.cuh>
#include <cub/detail/uninitialized_copy.cuh>
#include <cub/thread/thread_reduce.cuh>
#include <cub/thread/thread_scan.cuh>
#include <cub/util_ptx.cuh>
#include <cub/warp/warp_scan.cuh>
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 <typename T, int BlockDimX, int BlockDimY, int BlockDimZ, bool Memoize>
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<T, BLOCK_THREADS>;
/// 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<T, RAKING_THREADS>;
/// 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 <int ITERATION, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE T
GuardedReduce(T* raking_ptr, ScanOp scan_op, T raking_partial, constant_t<ITERATION> /*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<ITERATION + 1>);
}
/**
* @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 <typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE T
GuardedReduce(T* /*raking_ptr*/, ScanOp /*scan_op*/, T raking_partial, constant_t<SEGMENT_LENGTH> /*iteration*/)
{
return raking_partial;
}
/**
* @brief Templated copy
*
* @param out
* [out] Out array
*
* @param in
* [in] Input array
*/
template <int ITERATION>
_CCCL_DEVICE _CCCL_FORCEINLINE void CopySegment(T* out, T* in, constant_t<ITERATION> /*iteration*/)
{
out[ITERATION] = in[ITERATION];
CopySegment(out, in, constant_v<ITERATION + 1>);
}
/**
* @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<SEGMENT_LENGTH> /*iteration*/) {}
/// Performs upsweep raking reduction, returning the aggregate
template <typename ScanOp>
_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 <typename ScanOp>
_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 <typename ScanOp>
_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 <em>thread</em><sub>0</sub> 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 <typename ScanOp>
_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 <typename ScanOp>
_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 <em>thread</em><sub>0</sub> 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 <typename ScanOp>
_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 <typename ScanOp>
_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 <em>lane</em><sub>0</sub> 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
* <b>[<em>warp</em><sub>0</sub> only]</b> Call-back functor for specifying a thread
* block-wide prefix to be applied to all inputs.
*/
template <typename ScanOp, typename BlockPrefixCallbackOp>
_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 <typename ScanOp>
_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 <typename ScanOp>
_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 <em>lane</em><sub>0</sub> 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
* <b>[<em>warp</em><sub>0</sub> only]</b> Call-back functor for specifying a thread
* block-wide prefix to be applied to all inputs.
*/
template <typename ScanOp, typename BlockPrefixCallbackOp>
_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

514
cccl_upstream/cub/cub/block/specializations/block_scan_warp_scans.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/detail/uninitialized_copy.cuh>
#include <cub/util_ptx.cuh>
#include <cub/warp/warp_scan.cuh>
#include <cuda/__cmath/ceil_div.h>
#include <cuda/__ptx/instructions/get_sreg.h>
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 <typename T, int BlockDimX, int BlockDimY, int BlockDimZ>
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<T, WARP_THREADS>;
/// WarpScan utility type
using WarpAggregateScan = WarpScan<T, WARPS>;
/// 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 <typename ScanOp, int WARP>
_CCCL_DEVICE _CCCL_FORCEINLINE void
ApplyWarpAggregates(T& warp_prefix, ScanOp scan_op, T& block_aggregate, constant_t<WARP> /*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<WARP + 1>);
}
/**
* @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 <typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
ApplyWarpAggregates(T& /*warp_prefix*/, ScanOp /*scan_op*/, T& /*block_aggregate*/, constant_t<WARPS> /*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
* <b>[<em>lane</em><sub>WARP_THREADS - 1</sub> only]</b> Warp-wide aggregate reduction of
* input items
*
* @param[out] block_aggregate
* Threadblock-wide aggregate reduction of input items
*/
template <typename ScanOp>
_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
* <b>[<em>lane</em><sub>WARP_THREADS - 1</sub> only]</b> 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 <typename ScanOp>
_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 <em>thread</em><sub>0</sub> 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 <typename ScanOp>
_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 <typename ScanOp>
_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 <em>thread</em><sub>0</sub> 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 <typename ScanOp>
_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 <typename ScanOp>
_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 <em>lane</em><sub>0</sub> 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
* <b>[<em>warp</em><sub>0</sub> only]</b> Call-back functor for specifying a thread
* block-wide prefix to be applied to all inputs.
*/
template <typename ScanOp, typename BlockPrefixCallbackOp>
_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 <typename ScanOp>
_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 <typename ScanOp>
_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 <em>lane</em><sub>0</sub> 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
* <b>[<em>warp</em><sub>0</sub> only]</b> Call-back functor for specifying a thread
* block-wide prefix to be applied to all inputs.
*/
template <typename ScanOp, typename BlockPrefixCallbackOp>
_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

523
cccl_upstream/cub/cub/block/specializations/block_topk_air.cuh Normal file
View File

@@ -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 <cub/config.cuh>
#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
# pragma GCC system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
# pragma clang system_header
#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
# pragma system_header
#endif // no system header
#include <cub/block/block_scan.cuh>
#include <cub/block/radix_rank_sort_operations.cuh>
#include <cub/device/dispatch/dispatch_common.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <cuda/std/__bit/bit_cast.h>
#include <cuda/std/__type_traits/is_unsigned.h>
#include <cuda/std/cstdint>
CUB_NAMESPACE_BEGIN
namespace detail
{
template <typename SortKeyT>
struct compare_key_prefix_op
{
static_assert(::cuda::std::is_unsigned_v<SortKeyT>, "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 <typename KeyT, int ThreadsPerBlock, int ItemsPerThread, typename ValueT = NullType, int RadixBits = 8>
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<ValueT, NullType>;
using histo_counter_t = ::cuda::std::uint32_t;
using block_scan_t = BlockScan<histo_counter_t, threads_per_block, BLOCK_SCAN_WARP_SCANS>;
using traits = detail::radix::traits_t<KeyT>;
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<KeyT>;
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 <detail::topk::select SelectDirection, bool IsFullTile, typename DigitExtractorT, typename FilterOpT>
_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<int>(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 <typename detail::topk::select SelectDirection, bool IsFullTile, typename DecomposerT>
_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<bit_ordered_type>(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<bit_ordered_type>{prefix_mask, kth_key_prefix};
auto digit_extractor =
traits::template digit_extractor<fundamental_digit_extractor_t>(pass_begin_bit, pass_bits, decomposer);
compute_histograms<SelectDirection, IsFullTile>(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 <detail::topk::select SelectDirection, bool IsFullTile>
_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<bit_ordered_type(&)[ItemsPerThread]>(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<KeyT>)
{
const bit_ordered_type twiddled_minus_zero =
Traits<KeyT>::TwiddleIn(bit_ordered_type(1) << (8 * sizeof(bit_ordered_type) - 1));
const bit_ordered_type twiddled_zero = Traits<KeyT>::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<SelectDirection, IsFullTile>(
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<SelectDirection == detail::topk::select::min, ::cuda::std::less<>, ::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<KeyT>)
{
storage.stage.select.exchange.keys[selected_offset] =
(flip_back_bits[i / 32] & (1u << (i % 32))) ? KeyT(-0.0) : ::cuda::std::bit_cast<KeyT>(unsigned_keys[i]);
}
else
{
storage.stage.select.exchange.keys[selected_offset] = ::cuda::std::bit_cast<KeyT>(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<TempStorage_>
{};
_CCCL_DEVICE_API _CCCL_FORCEINLINE block_topk_air(TempStorage& storage)
: storage(storage.Alias())
, linear_tid(RowMajorTid(ThreadsPerBlock, 1, 1))
{}
template <detail::topk::select SelectDirection, bool IsFullTile>
_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<SelectDirection, IsFullTile>(keys, values, k, valid_items, begin_bit, end_bit);
}
template <detail::topk::select SelectDirection, bool IsFullTile>
_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<SelectDirection, IsFullTile>(keys, values, k, valid_items, begin_bit, end_bit);
}
};
} // namespace detail
CUB_NAMESPACE_END