8d75652949
Sources cloned and tree'd (no --depth):
- jd-opensource/xllm: ILU kernels, CUDA kernels, MoE kernels
- NVIDIA/cccl: CUB tuning/dispatch headers (block-level primitives)
- fla-org/flash-linear-attention: Triton GDN kernels
- NVIDIA/cutlass: grouped GEMM reference (read, not copied)
- Dao-AILab/flash-attention: attention kernel reference (SM80+, read only)
New CUDA kernels (from xllm, SM-agnostic, portable to BI-V100):
ex_engine/xllm_kernels/cuda/activation.cu (188 lines) — silu_and_mul, gelu
ex_engine/xllm_kernels/cuda/norm.cu (600 lines) — rms_norm, fused_add_rms_norm
ex_engine/xllm_kernels/cuda/rope.cu (258 lines) — rotary_embedding
ex_engine/xllm_kernels/cuda/block_copy.cu (209 lines) — copy_blocks, swap_blocks
ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu (101 lines) — KV cache ops
ex_engine/xllm_kernels/cuda/headers/ (5 headers for compilation)
ILU bridge kernel sources (from xllm, verified SAME as upstream):
ex_engine/xllm_kernels/ilu/ (10 files, 925 lines total)
— activation.cpp, attention.cpp, fused_moe.cpp, group_gemm.cpp,
matmul.cpp, norm.cpp, rope.cpp, ilu_ops_api.h, ixformer.h, utils.h
FLA Triton GDN kernels (for GatedDeltaNet without SM90+ FlashQLA):
ex_engine/fla_kernels/gated_delta_rule/ (7 files, 2370 lines)
— chunk_fwd.py (428), chunk.py (487), wy_fast.py (409),
fused_recurrent.py (392), naive.py (161), gate.py (380)
CCCL sync (12 tuning + 14 dispatch headers updated from NVIDIA/cccl):
cccl_upstream/cub/cub/device/dispatch/tuning/ — 12 changed files synced
cccl_upstream/cub/cub/device/dispatch/ — 14 changed dispatch files synced
Compilation targets for real machine (ivcore10):
1. CUDA kernels: --cuda-gpu-arch=ivcore10 via corex clang/16
2. ILU bridges: torch.utils.cpp_extension linking ixformer .so
3. FLA kernels: Triton JIT (if Triton works on BI-V100)
1482 lines
50 KiB
Plaintext
1482 lines
50 KiB
Plaintext
// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
|
|
// SPDX-FileCopyrightText: Copyright (c) 2011-2026, NVIDIA CORPORATION. All rights reserved.
|
|
// SPDX-License-Identifier: BSD-3
|
|
|
|
/**
|
|
* @file
|
|
* @brief cub::DeviceScan provides device-wide, parallel operations for
|
|
* computing a prefix scan across a sequence of data items residing
|
|
* within device-accessible memory.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include <cub/config.cuh>
|
|
|
|
#include <cub/util_namespace.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/agent/agent_scan.cuh>
|
|
#include <cub/detail/cc_dispatch.cuh>
|
|
#include <cub/detail/launcher/cuda_runtime.cuh>
|
|
#include <cub/detail/warpspeed/warpspeed.cuh>
|
|
#include <cub/device/dispatch/dispatch_common.cuh>
|
|
#include <cub/device/dispatch/kernels/kernel_scan.cuh>
|
|
#include <cub/device/dispatch/tuning/tuning_scan.cuh>
|
|
#include <cub/thread/thread_operators.cuh>
|
|
#include <cub/util_debug.cuh>
|
|
#include <cub/util_device.cuh>
|
|
#include <cub/util_math.cuh>
|
|
|
|
#include <thrust/system/cuda/detail/core/triple_chevron_launch.h>
|
|
#include <thrust/type_traits/unwrap_contiguous_iterator.h>
|
|
|
|
#include <cuda/__cmath/ceil_div.h>
|
|
#include <cuda/std/__algorithm/min.h>
|
|
#include <cuda/std/__functional/invoke.h>
|
|
#include <cuda/std/__host_stdlib/sstream>
|
|
#include <cuda/std/__iterator/readable_traits.h>
|
|
#include <cuda/std/__memory/construct_at.h>
|
|
#include <cuda/std/__type_traits/conditional.h>
|
|
#include <cuda/std/__type_traits/integral_constant.h>
|
|
#include <cuda/std/__type_traits/is_same.h>
|
|
#include <cuda/std/__type_traits/is_unsigned.h>
|
|
#include <cuda/std/__type_traits/void_t.h>
|
|
#include <cuda/std/__utility/move.h>
|
|
|
|
#include <cuda_runtime_api.h>
|
|
#include <cudaTypedefs.h>
|
|
|
|
CUB_NAMESPACE_BEGIN
|
|
|
|
namespace detail::scan
|
|
{
|
|
template <typename PolicySelector,
|
|
typename UnwrappedInputIteratorT,
|
|
typename UnwrappedOutputIteratorT,
|
|
typename ScanOpT,
|
|
typename InitValueT,
|
|
typename OffsetT,
|
|
typename AccumT,
|
|
ForceInclusive EnforceInclusive,
|
|
bool StableReductionOrder = false>
|
|
struct DeviceScanKernelSource
|
|
{
|
|
using ScanTileStateT = ScanTileState<AccumT>;
|
|
|
|
CUB_DEFINE_KERNEL_GETTER(
|
|
InitKernel,
|
|
DeviceScanInitKernel<PolicySelector, UnwrappedInputIteratorT, UnwrappedOutputIteratorT, ScanTileStateT, AccumT>)
|
|
|
|
CUB_DEFINE_KERNEL_GETTER(
|
|
ScanKernel,
|
|
DeviceScanKernel<PolicySelector,
|
|
UnwrappedInputIteratorT,
|
|
UnwrappedOutputIteratorT,
|
|
ScanTileStateT,
|
|
ScanOpT,
|
|
InitValueT,
|
|
OffsetT,
|
|
AccumT,
|
|
EnforceInclusive == ForceInclusive::Yes,
|
|
StableReductionOrder>)
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t InputSize()
|
|
{
|
|
return sizeof(it_value_t<UnwrappedInputIteratorT>);
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t InputAlign()
|
|
{
|
|
return alignof(it_value_t<UnwrappedInputIteratorT>);
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t OutputSize()
|
|
{
|
|
return sizeof(it_value_t<UnwrappedOutputIteratorT>);
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t OutputAlign()
|
|
{
|
|
return alignof(it_value_t<UnwrappedOutputIteratorT>);
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t AccumSize()
|
|
{
|
|
return sizeof(AccumT);
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t AccumAlign()
|
|
{
|
|
return alignof(AccumT);
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static ScanTileStateT TileState()
|
|
{
|
|
return {};
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t lookahead_tile_state_size()
|
|
{
|
|
return sizeof(warpspeed::tile_state_t<AccumT>);
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr ::cuda::std::size_t lookahead_tile_state_alignment()
|
|
{
|
|
return alignof(warpspeed::tile_state_t<AccumT>);
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr auto make_tile_state_kernel_arg(ScanTileStateT ts)
|
|
{
|
|
tile_state_kernel_arg_t<ScanTileStateT, AccumT> arg;
|
|
::cuda::std::__construct_at(&arg.lookback, ::cuda::std::move(ts));
|
|
return arg;
|
|
}
|
|
|
|
CUB_RUNTIME_FUNCTION static constexpr auto
|
|
lookahead_make_tile_state_kernel_arg(void* ts, ::cuda::std::uint32_t* atomic_counter = nullptr)
|
|
{
|
|
tile_state_kernel_arg_t<ScanTileStateT, AccumT> arg;
|
|
::cuda::std::__construct_at(
|
|
&arg.lookahead,
|
|
lookahead_tile_state_arg_t<AccumT>{static_cast<warpspeed::tile_state_t<AccumT>*>(ts), atomic_counter});
|
|
return arg;
|
|
}
|
|
};
|
|
|
|
// TODO(griwes): remove in CCCL 4.0 when we drop the scan dispatcher after publishing the tuning API
|
|
template <typename LegacyActivePolicy>
|
|
_CCCL_HOST_DEVICE_API constexpr auto convert_policy() -> ScanPolicy
|
|
{
|
|
// this does not convert any lookahead policy data, which is fine because we merged lookahead scan during the CCCL 3.4
|
|
// development cycle, so it never had user exposure through the policy_hub, and we can just only support it through
|
|
// the policy_selector.
|
|
using scan_policy_t = typename LegacyActivePolicy::ScanPolicyT;
|
|
return ScanPolicy{
|
|
ScanAlgorithm::lookback,
|
|
ScanLookbackPolicy{
|
|
scan_policy_t::BLOCK_THREADS,
|
|
scan_policy_t::ITEMS_PER_THREAD,
|
|
scan_policy_t::LOAD_ALGORITHM,
|
|
scan_policy_t::LOAD_MODIFIER,
|
|
scan_policy_t::STORE_ALGORITHM,
|
|
scan_policy_t::SCAN_ALGORITHM,
|
|
detail::lookback_delay_policy_from_type<typename scan_policy_t::detail::delay_constructor_t>},
|
|
ScanLookaheadPolicy{}};
|
|
}
|
|
|
|
// TODO(griwes): remove in CCCL 4.0 when we drop the scan dispatcher after publishing the tuning API
|
|
template <typename PolicyHub>
|
|
struct policy_selector_from_hub
|
|
{
|
|
[[nodiscard]] _CCCL_DEVICE_API constexpr auto operator()(::cuda::compute_capability /*cc*/) const -> ScanPolicy
|
|
{
|
|
return convert_policy<typename PolicyHub::MaxPolicy::ActivePolicy>();
|
|
}
|
|
};
|
|
} // namespace detail::scan
|
|
|
|
/******************************************************************************
|
|
* Dispatch
|
|
******************************************************************************/
|
|
|
|
/**
|
|
* @brief Utility class for dispatching the appropriately-tuned kernels for
|
|
* DeviceScan
|
|
*
|
|
* Deprecated [Since 3.5]
|
|
*
|
|
* @tparam InputIteratorT
|
|
* Random-access input iterator type for reading scan inputs @iterator
|
|
*
|
|
* @tparam OutputIteratorT
|
|
* Random-access output iterator type for writing scan outputs @iterator
|
|
*
|
|
* @tparam ScanOpT
|
|
* Binary scan functor type having member
|
|
* `auto operator()(const T &a, const U &b)`
|
|
*
|
|
* @tparam InitValueT
|
|
* The init_value element type for ScanOpT (cub::NullType for inclusive scans)
|
|
*
|
|
* @tparam OffsetT
|
|
* Unsigned integer type for global offsets
|
|
*
|
|
* @tparam EnforceInclusive
|
|
* Enum flag to specify whether to enforce inclusive scan.
|
|
*
|
|
*/
|
|
// TODO(griwes): Remove in CCCL 4.0
|
|
template <
|
|
typename InputIteratorT,
|
|
typename OutputIteratorT,
|
|
typename ScanOpT,
|
|
typename InitValueT,
|
|
typename OffsetT,
|
|
typename AccumT = ::cuda::std::__accumulator_t<ScanOpT,
|
|
cub::detail::it_value_t<InputIteratorT>,
|
|
::cuda::std::_If<::cuda::std::is_same_v<InitValueT, NullType>,
|
|
cub::detail::it_value_t<InputIteratorT>,
|
|
typename InitValueT::value_type>>,
|
|
ForceInclusive EnforceInclusive = ForceInclusive::No,
|
|
typename PolicyHub = detail::scan::
|
|
policy_hub<detail::it_value_t<InputIteratorT>, detail::it_value_t<OutputIteratorT>, AccumT, OffsetT, ScanOpT>,
|
|
typename KernelSource = detail::scan::DeviceScanKernelSource<
|
|
detail::scan::policy_selector_from_hub<PolicyHub>,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t<InputIteratorT>,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t<OutputIteratorT>,
|
|
ScanOpT,
|
|
InitValueT,
|
|
OffsetT,
|
|
AccumT,
|
|
EnforceInclusive>,
|
|
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
|
|
struct CCCL_DEPRECATED_BECAUSE("Use the tuning API for DeviceScan") DispatchScan
|
|
{
|
|
static_assert(::cuda::std::is_unsigned_v<OffsetT> && sizeof(OffsetT) >= 4,
|
|
"DispatchScan only supports unsigned offset types of at least 4-bytes");
|
|
|
|
//---------------------------------------------------------------------
|
|
// Constants and Types
|
|
//---------------------------------------------------------------------
|
|
|
|
static constexpr int INIT_KERNEL_THREADS = 128;
|
|
|
|
/// Device-accessible allocation of temporary storage. When nullptr, the
|
|
/// required allocation size is written to \p temp_storage_bytes and no work
|
|
/// is done.
|
|
void* d_temp_storage;
|
|
|
|
/// Reference to size in bytes of \p d_temp_storage allocation
|
|
size_t& temp_storage_bytes;
|
|
|
|
/// Iterator to the input sequence of data items
|
|
InputIteratorT d_in;
|
|
|
|
/// Iterator to the output sequence of data items
|
|
OutputIteratorT d_out;
|
|
|
|
/// Binary scan functor
|
|
ScanOpT scan_op;
|
|
|
|
/// Initial value to seed the exclusive scan
|
|
InitValueT init_value;
|
|
|
|
/// Total number of input items (i.e., the length of \p d_in)
|
|
OffsetT num_items;
|
|
|
|
/// CUDA stream to launch kernels within. Default is stream<sub>0</sub>.
|
|
cudaStream_t stream;
|
|
|
|
int ptx_version;
|
|
|
|
KernelSource kernel_source;
|
|
|
|
KernelLauncherFactory launcher_factory;
|
|
|
|
/**
|
|
*
|
|
* @param[in] d_temp_storage
|
|
* Device-accessible allocation of temporary storage. When `nullptr`, the
|
|
* required allocation size is written to `temp_storage_bytes` and no
|
|
* work is done.
|
|
*
|
|
* @param[in,out] temp_storage_bytes
|
|
* Reference to size in bytes of `d_temp_storage` allocation
|
|
*
|
|
* @param[in] d_in
|
|
* Iterator to the input sequence of data items
|
|
*
|
|
* @param[out] d_out
|
|
* Iterator to the output sequence of data items
|
|
*
|
|
* @param[in] num_items
|
|
* Total number of input items (i.e., the length of `d_in`)
|
|
*
|
|
* @param[in] scan_op
|
|
* Binary scan functor
|
|
*
|
|
* @param[in] init_value
|
|
* Initial value to seed the exclusive scan
|
|
*
|
|
* @param[in] stream
|
|
* **[optional]** CUDA stream to launch kernels within.
|
|
* Default is stream<sub>0</sub>.
|
|
*
|
|
* @param[in] kernel_source
|
|
* Object specifying implementation kernels
|
|
*
|
|
* @param[in] launcher_factory
|
|
* Object to execute implementation kernels on the given stream
|
|
*/
|
|
// TODO(griwes): Remove in CCCL 4.0
|
|
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE DispatchScan(
|
|
void* d_temp_storage,
|
|
size_t& temp_storage_bytes,
|
|
InputIteratorT d_in,
|
|
OutputIteratorT d_out,
|
|
OffsetT num_items,
|
|
ScanOpT scan_op,
|
|
InitValueT init_value,
|
|
cudaStream_t stream,
|
|
int ptx_version,
|
|
KernelSource kernel_source = {},
|
|
KernelLauncherFactory launcher_factory = {})
|
|
: d_temp_storage(d_temp_storage)
|
|
, temp_storage_bytes(temp_storage_bytes)
|
|
, d_in(d_in)
|
|
, d_out(d_out)
|
|
, scan_op(scan_op)
|
|
, init_value(init_value)
|
|
, num_items(num_items)
|
|
, stream(stream)
|
|
, ptx_version(ptx_version)
|
|
, kernel_source(kernel_source)
|
|
, launcher_factory(launcher_factory)
|
|
{}
|
|
|
|
template <typename ActivePolicyT, typename InitKernelT, typename ScanKernelT>
|
|
CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t
|
|
Invoke(InitKernelT init_kernel, ScanKernelT scan_kernel, ActivePolicyT policy = {})
|
|
{
|
|
// `LOAD_LDG` makes in-place execution UB and doesn't lead to better
|
|
// performance.
|
|
policy.CheckLoadModifier();
|
|
|
|
// Number of input tiles
|
|
const int tile_size = policy.Scan().ThreadsPerBlock() * policy.Scan().ItemsPerThread();
|
|
const int num_tiles = static_cast<int>(::cuda::ceil_div(num_items, tile_size));
|
|
|
|
auto tile_state = kernel_source.TileState();
|
|
|
|
// Specify temporary storage allocation requirements
|
|
size_t allocation_sizes[1];
|
|
if (const auto error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0])))
|
|
{
|
|
return error; // bytes needed for tile status descriptors
|
|
}
|
|
|
|
// Compute allocation pointers into the single storage blob (or compute
|
|
// the necessary size of the blob)
|
|
void* allocations[1] = {};
|
|
if (const auto error =
|
|
CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Return if the caller is simply requesting the size of the storage allocation, or the problem is empty
|
|
if (d_temp_storage == nullptr || num_items == 0)
|
|
{
|
|
return cudaSuccess;
|
|
}
|
|
|
|
// Construct the tile status interface
|
|
if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0])))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Log init_kernel configuration
|
|
const int init_grid_size = ::cuda::ceil_div(num_tiles, INIT_KERNEL_THREADS);
|
|
|
|
#ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, INIT_KERNEL_THREADS, (long long) stream);
|
|
#endif // CUB_DEBUG_LOG
|
|
|
|
// Invoke init_kernel to initialize tile descriptors
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(init_grid_size, INIT_KERNEL_THREADS, 0, stream, /* dependent_launch */ ptx_version >= 900)
|
|
.doit(init_kernel, kernel_source.make_tile_state_kernel_arg(tile_state), num_tiles)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Get SM occupancy for scan_kernel
|
|
int scan_sm_occupancy;
|
|
if (const auto error =
|
|
CubDebug(launcher_factory.MaxSmOccupancy(scan_sm_occupancy, scan_kernel, policy.Scan().ThreadsPerBlock())))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Get max x-dimension of grid
|
|
int max_dim_x;
|
|
if (const auto error = CubDebug(launcher_factory.MaxGridDimX(max_dim_x)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Run grids in epochs (in case number of tiles exceeds max x-dimension
|
|
const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x);
|
|
for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size)
|
|
{
|
|
// Log scan_kernel configuration
|
|
#ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking %d scan_kernel<<<%d, %d, 0, %lld>>>(), %d items "
|
|
"per thread, %d SM occupancy\n",
|
|
start_tile,
|
|
scan_grid_size,
|
|
policy.Scan().ThreadsPerBlock(),
|
|
(long long) stream,
|
|
policy.Scan().ItemsPerThread(),
|
|
scan_sm_occupancy);
|
|
#endif // CUB_DEBUG_LOG
|
|
|
|
// Invoke scan_kernel
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(
|
|
scan_grid_size, policy.Scan().ThreadsPerBlock(), 0, stream, /* dependent_launch */ ptx_version >= 900)
|
|
.doit(scan_kernel,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in),
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out),
|
|
kernel_source.make_tile_state_kernel_arg(tile_state),
|
|
start_tile,
|
|
scan_op,
|
|
init_value,
|
|
num_items,
|
|
/* num_stages, unused */ 1)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
}
|
|
|
|
return cudaSuccess;
|
|
}
|
|
|
|
// do check in separate function, so error message contains the required SMEM in the error novel
|
|
template <int SMemSizeForSingleStage>
|
|
CUB_RUNTIME_FUNCTION static void __check_smem()
|
|
{
|
|
static_assert(SMemSizeForSingleStage <= detail::max_smem_per_block,
|
|
"Single-stage lookahead scan exceeds architecture independent SMEM (48KiB)");
|
|
}
|
|
|
|
template <typename PolicyGetter>
|
|
CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t __invoke_lookahead_algorithm(PolicyGetter policy_getter)
|
|
{
|
|
#if __cccl_ptx_isa >= 860
|
|
if (num_items == 0)
|
|
{
|
|
temp_storage_bytes = 1; // just fulfill the contract that CUB always requires some temporary storage
|
|
return cudaSuccess;
|
|
}
|
|
|
|
CUB_DETAIL_CONSTEXPR_ISH const ScanLookaheadPolicy lookahead_policy = policy_getter().lookahead;
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(lookahead_policy.reduce_and_scan_warps >= 1,
|
|
"Lookahead scan policy have at least 1 warp for reducing and scanning");
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(
|
|
lookahead_policy.items_per_thread >= 1, "Lookahead scan policy must have at least 1 item per thread");
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(lookahead_policy.lookahead_items_per_thread >= 1,
|
|
"Lookahead scan policy must look ahead at least 1 item per thread");
|
|
|
|
const int grid_dim =
|
|
static_cast<int>(::cuda::ceil_div(num_items, static_cast<OffsetT>(lookahead_policy.tile_size())));
|
|
|
|
if (d_temp_storage == nullptr)
|
|
{
|
|
temp_storage_bytes = static_cast<size_t>(grid_dim) * kernel_source.lookahead_tile_state_size();
|
|
return cudaSuccess;
|
|
}
|
|
|
|
if (num_items == 0)
|
|
{
|
|
return cudaSuccess;
|
|
}
|
|
|
|
int sm_count = 0;
|
|
if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count)))
|
|
{
|
|
return error;
|
|
}
|
|
// Maximum dynamic shared memory size that we can use for temporary storage.
|
|
int max_dynamic_smem_size{};
|
|
if (const auto error =
|
|
CubDebug(launcher_factory.max_dynamic_smem_size_for(max_dynamic_smem_size, kernel_source.ScanKernel())))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// TODO(bgruber): we probably need to ensure alignment of d_temp_storage
|
|
_CCCL_ASSERT(::cuda::is_aligned(d_temp_storage, kernel_source.lookahead_tile_state_alignment()), "");
|
|
|
|
auto scan_kernel = kernel_source.ScanKernel();
|
|
[[maybe_unused]] auto kernel_src = kernel_source; // need to pull a copy to not access `this` during const. eval.
|
|
CUB_DETAIL_CONSTEXPR_ISH int smem_size_1_stage = detail::scan::smem_for_stages(
|
|
lookahead_policy,
|
|
1,
|
|
static_cast<int>(kernel_src.InputSize()),
|
|
static_cast<int>(kernel_src.InputAlign()),
|
|
static_cast<int>(kernel_src.OutputAlign()),
|
|
static_cast<int>(kernel_src.AccumSize()),
|
|
static_cast<int>(kernel_src.AccumAlign()));
|
|
# if defined(CUB_DEFINE_RUNTIME_POLICIES)
|
|
_CCCL_ASSERT(smem_size_1_stage <= int{detail::max_smem_per_block},
|
|
"Single-stage lookahead scan exceeds architecture independent SMEM (48KiB)");
|
|
# else // defined(CUB_DEFINE_RUNTIME_POLICIES)
|
|
__check_smem<smem_size_1_stage>();
|
|
# endif // defined(CUB_DEFINE_RUNTIME_POLICIES)
|
|
|
|
int num_stages = 1;
|
|
int smem_size = smem_size_1_stage;
|
|
|
|
// When launched from the host, maximize the number of stages that we can fit inside the shared memory.
|
|
NV_IF_TARGET(
|
|
NV_IS_HOST, ({
|
|
// number of stages to have an even workload across all SMs (improves small problem sizes), assuming
|
|
// 1 CTA per SM +1 since it tends to improve performance
|
|
// TODO(bgruber): make the +1 a tuning parameter
|
|
const int max_stages_for_even_workload = static_cast<int>(
|
|
::cuda::ceil_div(num_items, static_cast<OffsetT>(sm_count * lookahead_policy.tile_size())) + 1);
|
|
|
|
while (num_stages <= max_stages_for_even_workload)
|
|
{
|
|
const int next_smem_size = detail::scan::smem_for_stages(
|
|
lookahead_policy,
|
|
num_stages + 1,
|
|
static_cast<int>(kernel_source.InputSize()),
|
|
static_cast<int>(kernel_source.InputAlign()),
|
|
static_cast<int>(kernel_source.OutputAlign()),
|
|
static_cast<int>(kernel_source.AccumSize()),
|
|
static_cast<int>(kernel_source.AccumAlign()));
|
|
if (next_smem_size > max_dynamic_smem_size)
|
|
{
|
|
// This number of stages failed, so stay at the current settings
|
|
break;
|
|
}
|
|
|
|
smem_size = next_smem_size;
|
|
++num_stages;
|
|
}
|
|
|
|
// Set scan kernel's max shared memory limit to the max smem value. We might not use all of it, but it prevents
|
|
// multiple kernels from overwriting the max shared memory limit by different values.
|
|
//
|
|
// TODO: Since CTK 13.2 we can use CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE to allow non-portable shared memory
|
|
// sizes, however we need something that works even with older CTKs.
|
|
if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, max_dynamic_smem_size))
|
|
{
|
|
return error;
|
|
}
|
|
}))
|
|
|
|
// Invoke init kernel
|
|
{
|
|
constexpr auto init_kernel_threads = 128;
|
|
const auto init_grid_size = ::cuda::ceil_div(grid_dim, init_kernel_threads);
|
|
|
|
# ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking DeviceScanInitKernel<<<%d, %d, 0, %lld>>>()\n",
|
|
init_grid_size,
|
|
init_kernel_threads,
|
|
(long long) stream);
|
|
# endif // CUB_DEBUG_LOG
|
|
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(init_grid_size,
|
|
init_kernel_threads,
|
|
0,
|
|
stream,
|
|
/* dependent_launch */ ptx_version >= 900)
|
|
.doit(kernel_source.InitKernel(),
|
|
kernel_source.lookahead_make_tile_state_kernel_arg(d_temp_storage),
|
|
grid_dim)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
}
|
|
|
|
// Invoke scan kernel
|
|
{
|
|
const int block_dim = detail::scan::num_total_threads(lookahead_policy);
|
|
|
|
# ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking DeviceScanKernel<<<%d, %d, %d, %lld>>>()\n", grid_dim, block_dim, smem_size, (long long) stream);
|
|
# endif // CUB_DEBUG_LOG
|
|
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(grid_dim, block_dim, smem_size, stream, /* dependent_launch */ ptx_version >= 900)
|
|
.doit(scan_kernel,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in),
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out),
|
|
kernel_source.lookahead_make_tile_state_kernel_arg(d_temp_storage),
|
|
/* start_tile, unused */ 0,
|
|
::cuda::std::move(scan_op),
|
|
init_value,
|
|
num_items,
|
|
num_stages)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
}
|
|
#else // __cccl_ptx_isa >= 860
|
|
static_assert(sizeof(policy_getter) == 0,
|
|
"Implementation bug: Tuning policy selected lookahead, but supported PTX ISA is too low");
|
|
#endif // __cccl_ptx_isa >= 860
|
|
return cudaSuccess;
|
|
}
|
|
|
|
template <typename PolicyGetter>
|
|
CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t __invoke_lookback_algorithm(PolicyGetter policy_getter)
|
|
{
|
|
CUB_DETAIL_CONSTEXPR_ISH const ScanLookbackPolicy active_policy = policy_getter().lookback;
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(
|
|
active_policy.threads_per_block >= 1, "Lookback scan policy must have at least 1 thread per block");
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(
|
|
active_policy.items_per_thread >= 1, "Lookback scan policy must have at least 1 item per thread");
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(active_policy.load_modifier != CacheLoadModifier::LOAD_LDG,
|
|
"The memory consistency model does not apply to texture accesses");
|
|
|
|
// Number of input tiles
|
|
const int tile_size = active_policy.threads_per_block * active_policy.items_per_thread;
|
|
const int num_tiles = static_cast<int>(::cuda::ceil_div(num_items, tile_size));
|
|
|
|
auto tile_state = kernel_source.TileState();
|
|
|
|
// Specify temporary storage allocation requirements
|
|
size_t allocation_sizes[1];
|
|
if (const auto error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0])))
|
|
{
|
|
return error; // bytes needed for tile status descriptors
|
|
}
|
|
|
|
// Compute allocation pointers into the single storage blob (or compute
|
|
// the necessary size of the blob)
|
|
void* allocations[1] = {};
|
|
if (const auto error =
|
|
CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Return if the caller is simply requesting the size of the storage allocation, or the problem is empty
|
|
if (d_temp_storage == nullptr || num_items == 0)
|
|
{
|
|
return cudaSuccess;
|
|
}
|
|
|
|
// Construct the tile status interface
|
|
if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0])))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Log init_kernel configuration
|
|
constexpr int init_kernel_threads = 128;
|
|
const int init_grid_size = ::cuda::ceil_div(num_tiles, init_kernel_threads);
|
|
|
|
#ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, init_kernel_threads, (long long) stream);
|
|
#endif // CUB_DEBUG_LOG
|
|
|
|
// Invoke init_kernel to initialize tile descriptors
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(init_grid_size, init_kernel_threads, 0, stream, /* dependent_launch */ ptx_version >= 900)
|
|
.doit(kernel_source.InitKernel(), kernel_source.make_tile_state_kernel_arg(tile_state), num_tiles)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Get SM occupancy for scan_kernel
|
|
int scan_sm_occupancy;
|
|
if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy(
|
|
scan_sm_occupancy, kernel_source.ScanKernel(), active_policy.threads_per_block)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Get max x-dimension of grid
|
|
int max_dim_x;
|
|
if (const auto error = CubDebug(launcher_factory.MaxGridDimX(max_dim_x)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Run grids in epochs (in case number of tiles exceeds max x-dimension
|
|
const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x);
|
|
for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size)
|
|
{
|
|
// Log scan_kernel configuration
|
|
#ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking %d scan_kernel<<<%d, %d, 0, %lld>>>(), %d items "
|
|
"per thread, %d SM occupancy\n",
|
|
start_tile,
|
|
scan_grid_size,
|
|
active_policy.threads_per_block,
|
|
(long long) stream,
|
|
active_policy.items_per_thread,
|
|
scan_sm_occupancy);
|
|
#endif // CUB_DEBUG_LOG
|
|
|
|
// Invoke scan_kernel
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(
|
|
scan_grid_size, active_policy.threads_per_block, 0, stream, /* dependent_launch */ ptx_version >= 900)
|
|
.doit(kernel_source.ScanKernel(),
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in),
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out),
|
|
kernel_source.make_tile_state_kernel_arg(tile_state),
|
|
start_tile,
|
|
scan_op,
|
|
init_value,
|
|
num_items,
|
|
/* num_stages, unused */ 1)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
}
|
|
|
|
return cudaSuccess;
|
|
}
|
|
|
|
template <typename ActivePolicyT>
|
|
CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t Invoke(ActivePolicyT = {})
|
|
{
|
|
struct policy_getter
|
|
{
|
|
// host-device not api, because clang warns about exclude_from_explicit_instantiation in local types
|
|
_CCCL_HOST_DEVICE _CCCL_FORCEINLINE constexpr auto operator()() const
|
|
{
|
|
return detail::scan::convert_policy<ActivePolicyT>();
|
|
}
|
|
};
|
|
|
|
if CUB_DETAIL_CONSTEXPR_ISH (policy_getter{}().algorithm == ScanAlgorithm::lookahead)
|
|
{
|
|
return __invoke_lookahead_algorithm(policy_getter{});
|
|
}
|
|
else
|
|
{
|
|
return __invoke_lookback_algorithm(policy_getter{});
|
|
}
|
|
}
|
|
|
|
/**
|
|
* @brief Internal dispatch routine
|
|
*
|
|
* @param[in] d_temp_storage
|
|
* Device-accessible allocation of temporary storage. When `nullptr`, the
|
|
* required allocation size is written to `temp_storage_bytes` and no
|
|
* work is done.
|
|
*
|
|
* @param[in,out] temp_storage_bytes
|
|
* Reference to size in bytes of `d_temp_storage` allocation
|
|
*
|
|
* @param[in] d_in
|
|
* Iterator to the input sequence of data items
|
|
*
|
|
* @param[out] d_out
|
|
* Iterator to the output sequence of data items
|
|
*
|
|
* @param[in] scan_op
|
|
* Binary scan functor
|
|
*
|
|
* @param[in] init_value
|
|
* Initial value to seed the exclusive scan
|
|
*
|
|
* @param[in] num_items
|
|
* Total number of input items (i.e., the length of `d_in`)
|
|
*
|
|
* @param[in] stream
|
|
* **[optional]** CUDA stream to launch kernels within.
|
|
* Default is stream<sub>0</sub>.
|
|
*
|
|
* @param[in] kernel_source
|
|
* Object specifying implementation kernels
|
|
*
|
|
* @param[in] launcher_factory
|
|
* Object to execute implementation kernels on the given stream
|
|
*
|
|
* @param[in] max_policy
|
|
* Struct encoding chain of algorithm tuning policies
|
|
*/
|
|
// TODO(griwes): Remove in CCCL 4.0
|
|
template <typename MaxPolicyT = typename PolicyHub::MaxPolicy>
|
|
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE static cudaError_t Dispatch(
|
|
void* d_temp_storage,
|
|
size_t& temp_storage_bytes,
|
|
InputIteratorT d_in,
|
|
OutputIteratorT d_out,
|
|
ScanOpT scan_op,
|
|
InitValueT init_value,
|
|
OffsetT num_items,
|
|
cudaStream_t stream,
|
|
KernelSource kernel_source = {},
|
|
KernelLauncherFactory launcher_factory = {},
|
|
MaxPolicyT max_policy = {})
|
|
{
|
|
// Get PTX version
|
|
int ptx_version = 0;
|
|
if (const auto error = CubDebug(launcher_factory.PtxVersion(ptx_version)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Create dispatch functor
|
|
DispatchScan dispatch(
|
|
d_temp_storage,
|
|
temp_storage_bytes,
|
|
d_in,
|
|
d_out,
|
|
num_items,
|
|
scan_op,
|
|
init_value,
|
|
stream,
|
|
ptx_version,
|
|
kernel_source,
|
|
launcher_factory);
|
|
|
|
// Dispatch to chained policy
|
|
return CubDebug(max_policy.Invoke(ptx_version, dispatch));
|
|
}
|
|
};
|
|
|
|
namespace detail::scan
|
|
{
|
|
// do check in separate function, so error message contains the required SMEM in the error novel
|
|
template <int SMemSizeForSingleStage>
|
|
CUB_RUNTIME_FUNCTION void check_lookahead_smem()
|
|
{
|
|
static_assert(SMemSizeForSingleStage <= detail::max_smem_per_block,
|
|
"Single-stage lookahead scan exceeds architecture independent SMEM (48KiB)");
|
|
}
|
|
|
|
template <typename PolicyGetter,
|
|
typename InputIteratorT,
|
|
typename OutputIteratorT,
|
|
typename ScanOpT,
|
|
typename InitValueT,
|
|
typename OffsetT,
|
|
typename KernelSource,
|
|
typename KernelLauncherFactory>
|
|
CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t invoke_lookback(
|
|
PolicyGetter policy_getter,
|
|
void* d_temp_storage,
|
|
size_t& temp_storage_bytes,
|
|
InputIteratorT d_in,
|
|
OutputIteratorT d_out,
|
|
ScanOpT scan_op,
|
|
InitValueT init_value,
|
|
OffsetT num_items,
|
|
cudaStream_t stream,
|
|
bool dependent_launch,
|
|
KernelSource kernel_source,
|
|
KernelLauncherFactory launcher_factory)
|
|
{
|
|
CUB_DETAIL_CONSTEXPR_ISH const ScanLookbackPolicy active_policy = policy_getter().lookback;
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(
|
|
active_policy.threads_per_block >= 1, "Lookback scan policy must have at least 1 thread per block");
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(
|
|
active_policy.items_per_thread >= 1, "Lookback scan policy must have at least 1 item per thread");
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(active_policy.load_modifier != CacheLoadModifier::LOAD_LDG,
|
|
"The memory consistency model does not apply to texture accesses");
|
|
|
|
// Number of input tiles
|
|
const int tile_size = active_policy.threads_per_block * active_policy.items_per_thread;
|
|
const int num_tiles = static_cast<int>(::cuda::ceil_div(num_items, tile_size));
|
|
|
|
auto tile_state = kernel_source.TileState();
|
|
|
|
// Specify temporary storage allocation requirements
|
|
size_t allocation_sizes[1];
|
|
if (const auto error = CubDebug(tile_state.AllocationSize(num_tiles, allocation_sizes[0])))
|
|
{
|
|
return error; // bytes needed for tile status descriptors
|
|
}
|
|
|
|
// Compute allocation pointers into the single storage blob (or compute
|
|
// the necessary size of the blob)
|
|
void* allocations[1] = {};
|
|
if (const auto error =
|
|
CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Return if the caller is simply requesting the size of the storage allocation, or the problem is empty
|
|
if (d_temp_storage == nullptr || num_items == 0)
|
|
{
|
|
return cudaSuccess;
|
|
}
|
|
|
|
// Construct the tile status interface
|
|
if (const auto error = CubDebug(tile_state.Init(num_tiles, allocations[0], allocation_sizes[0])))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Log init_kernel configuration
|
|
constexpr int init_kernel_threads = 128;
|
|
const int init_grid_size = ::cuda::ceil_div(num_tiles, init_kernel_threads);
|
|
|
|
#ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking init_kernel<<<%d, %d, 0, %lld>>>()\n", init_grid_size, init_kernel_threads, (long long) stream);
|
|
#endif // CUB_DEBUG_LOG
|
|
|
|
// Invoke init_kernel to initialize tile descriptors
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(init_grid_size, init_kernel_threads, 0, stream, dependent_launch)
|
|
.doit(kernel_source.InitKernel(), kernel_source.make_tile_state_kernel_arg(tile_state), num_tiles)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Get SM occupancy for scan_kernel
|
|
int scan_sm_occupancy;
|
|
if (const auto error = CubDebug(launcher_factory.MaxSmOccupancy(
|
|
scan_sm_occupancy, kernel_source.ScanKernel(), active_policy.threads_per_block)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Get max x-dimension of grid
|
|
int max_dim_x;
|
|
if (const auto error = CubDebug(launcher_factory.MaxGridDimX(max_dim_x)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Run grids in epochs (in case number of tiles exceeds max x-dimension
|
|
const int scan_grid_size = ::cuda::std::min(num_tiles, max_dim_x);
|
|
for (int start_tile = 0; start_tile < num_tiles; start_tile += scan_grid_size)
|
|
{
|
|
// Log scan_kernel configuration
|
|
#ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking %d scan_kernel<<<%d, %d, 0, %lld>>>(), %d items "
|
|
"per thread, %d SM occupancy\n",
|
|
start_tile,
|
|
scan_grid_size,
|
|
active_policy.threads_per_block,
|
|
(long long) stream,
|
|
active_policy.items_per_thread,
|
|
scan_sm_occupancy);
|
|
#endif // CUB_DEBUG_LOG
|
|
|
|
// Invoke scan_kernel
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(scan_grid_size, active_policy.threads_per_block, 0, stream, dependent_launch)
|
|
.doit(kernel_source.ScanKernel(),
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in),
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out),
|
|
kernel_source.make_tile_state_kernel_arg(tile_state),
|
|
start_tile,
|
|
scan_op,
|
|
init_value,
|
|
num_items,
|
|
/* num_stages, unused */ 1)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
}
|
|
|
|
return cudaSuccess;
|
|
}
|
|
|
|
template <typename PolicyGetter,
|
|
typename InputIteratorT,
|
|
typename OutputIteratorT,
|
|
typename ScanOpT,
|
|
typename InitValueT,
|
|
typename OffsetT,
|
|
typename KernelSource,
|
|
typename KernelLauncherFactory>
|
|
CUB_RUNTIME_FUNCTION _CCCL_HOST _CCCL_FORCEINLINE cudaError_t invoke_lookahead(
|
|
PolicyGetter policy_getter,
|
|
void* d_temp_storage,
|
|
size_t& temp_storage_bytes,
|
|
InputIteratorT d_in,
|
|
OutputIteratorT d_out,
|
|
ScanOpT scan_op,
|
|
InitValueT init_value,
|
|
OffsetT num_items,
|
|
cudaStream_t stream,
|
|
bool dependent_launch,
|
|
bool atomic_scheduling,
|
|
KernelSource kernel_source,
|
|
KernelLauncherFactory launcher_factory)
|
|
{
|
|
#if __cccl_ptx_isa >= 860
|
|
if (num_items == 0)
|
|
{
|
|
temp_storage_bytes = 1; // just fulfill the contract that CUB always requires some temporary storage
|
|
return cudaSuccess;
|
|
}
|
|
|
|
CUB_DETAIL_CONSTEXPR_ISH const ScanLookaheadPolicy lookahead_policy = policy_getter().lookahead;
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(lookahead_policy.reduce_and_scan_warps >= 1,
|
|
"Lookahead scan policy have at least 1 warp for reducing and scanning");
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(
|
|
lookahead_policy.items_per_thread >= 1, "Lookahead scan policy must have at least 1 item per thread");
|
|
CUB_DETAIL_STATIC_ISH_ASSERT(lookahead_policy.lookahead_items_per_thread >= 1,
|
|
"Lookahead scan policy must look ahead at least 1 item per thread");
|
|
|
|
const int num_tiles =
|
|
static_cast<int>(::cuda::ceil_div(num_items, static_cast<OffsetT>(lookahead_policy.tile_size())));
|
|
|
|
size_t allocation_sizes[2] = {
|
|
static_cast<size_t>(num_tiles) * kernel_source.lookahead_tile_state_size(), sizeof(::cuda::std::uint32_t)};
|
|
void* allocations[2] = {};
|
|
if (const auto error =
|
|
CubDebug(detail::alias_temporaries(d_temp_storage, temp_storage_bytes, allocations, allocation_sizes)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
if (d_temp_storage == nullptr)
|
|
{
|
|
return cudaSuccess;
|
|
}
|
|
|
|
void* d_tile_state = allocations[0];
|
|
::cuda::std::uint32_t* d_atomic_counter = static_cast<::cuda::std::uint32_t*>(allocations[1]);
|
|
|
|
int sm_count = 0;
|
|
if (const auto error = CubDebug(launcher_factory.MultiProcessorCount(sm_count)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
const int scan_grid_dim = atomic_scheduling ? (::cuda::std::min) (sm_count, num_tiles) : num_tiles;
|
|
// Maximum dynamic shared memory size that we can use for temporary storage.
|
|
int max_dynamic_smem_size{};
|
|
if (const auto error =
|
|
CubDebug(launcher_factory.max_dynamic_smem_size_for(max_dynamic_smem_size, kernel_source.ScanKernel())))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
auto scan_kernel = kernel_source.ScanKernel();
|
|
[[maybe_unused]] auto kernel_src = kernel_source; // need to pull a copy to not access `this` during const. eval.
|
|
CUB_DETAIL_CONSTEXPR_ISH int smem_size_1_stage = detail::scan::smem_for_stages(
|
|
lookahead_policy,
|
|
1,
|
|
static_cast<int>(kernel_src.InputSize()),
|
|
static_cast<int>(kernel_src.InputAlign()),
|
|
static_cast<int>(kernel_src.OutputAlign()),
|
|
static_cast<int>(kernel_src.AccumSize()),
|
|
static_cast<int>(kernel_src.AccumAlign()));
|
|
# if defined(CUB_DEFINE_RUNTIME_POLICIES)
|
|
_CCCL_ASSERT(smem_size_1_stage <= int{detail::max_smem_per_block},
|
|
"Single-stage lookahead scan exceeds architecture independent SMEM (48KiB)");
|
|
# else // defined(CUB_DEFINE_RUNTIME_POLICIES)
|
|
check_lookahead_smem<smem_size_1_stage>();
|
|
# endif // defined(CUB_DEFINE_RUNTIME_POLICIES)
|
|
|
|
int num_stages = 1;
|
|
int smem_size = smem_size_1_stage;
|
|
|
|
// When launched from the host, maximize the number of stages that we can fit inside the shared memory.
|
|
NV_IF_TARGET(
|
|
NV_IS_HOST, ({
|
|
// number of stages to have an even workload across all SMs (improves small problem sizes), assuming
|
|
// 1 CTA per SM +1 since it tends to improve performance
|
|
// TODO(bgruber): make the +1 a tuning parameter
|
|
const int max_stages_for_even_workload = static_cast<int>(
|
|
::cuda::ceil_div(num_items, static_cast<OffsetT>(sm_count * lookahead_policy.tile_size())) + 1);
|
|
|
|
while (num_stages <= max_stages_for_even_workload)
|
|
{
|
|
const int next_smem_size = detail::scan::smem_for_stages(
|
|
lookahead_policy,
|
|
num_stages + 1,
|
|
static_cast<int>(kernel_source.InputSize()),
|
|
static_cast<int>(kernel_source.InputAlign()),
|
|
static_cast<int>(kernel_source.OutputAlign()),
|
|
static_cast<int>(kernel_source.AccumSize()),
|
|
static_cast<int>(kernel_source.AccumAlign()));
|
|
if (next_smem_size > max_dynamic_smem_size)
|
|
{
|
|
// This number of stages failed, so stay at the current settings
|
|
break;
|
|
}
|
|
|
|
smem_size = next_smem_size;
|
|
++num_stages;
|
|
}
|
|
|
|
// Set scan kernel's max shared memory limit to the max smem value. We might not use all of it, but it prevents
|
|
// multiple kernels from overwriting the max shared memory limit by different values.
|
|
//
|
|
// TODO: Since CTK 13.2 we can use CU_LAUNCH_ATTRIBUTE_SHARED_MEMORY_MODE to allow non-portable shared memory
|
|
// sizes, however we need something that works even with older CTKs.
|
|
if (const auto error = launcher_factory.set_max_dynamic_smem_size_for(scan_kernel, max_dynamic_smem_size))
|
|
{
|
|
return error;
|
|
}
|
|
}))
|
|
|
|
// Invoke init kernel
|
|
{
|
|
constexpr auto init_kernel_threads = 128;
|
|
const auto init_grid_size = ::cuda::ceil_div(num_tiles, init_kernel_threads);
|
|
|
|
# ifdef CUB_DEBUG_LOG
|
|
_CubLog("Invoking DeviceScanInitKernel<<<%d, %d, 0, %lld>>>()\n",
|
|
init_grid_size,
|
|
init_kernel_threads,
|
|
(long long) stream);
|
|
# endif // CUB_DEBUG_LOG
|
|
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(init_grid_size, init_kernel_threads, 0, stream, dependent_launch)
|
|
.doit(kernel_source.InitKernel(),
|
|
kernel_source.lookahead_make_tile_state_kernel_arg(d_tile_state, d_atomic_counter),
|
|
num_tiles)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
}
|
|
|
|
// Invoke scan kernel
|
|
{
|
|
const int block_dim = detail::scan::num_total_threads(lookahead_policy);
|
|
# ifdef CUB_DEBUG_LOG
|
|
_CubLog(
|
|
"Invoking DeviceScanKernel<<<%d, %d, %d, %lld>>>()\n", scan_grid_dim, block_dim, smem_size, (long long) stream);
|
|
# endif // CUB_DEBUG_LOG
|
|
|
|
if (const auto error = CubDebug(
|
|
launcher_factory(scan_grid_dim, block_dim, smem_size, stream, dependent_launch)
|
|
.doit(scan_kernel,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_in),
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(d_out),
|
|
kernel_source.lookahead_make_tile_state_kernel_arg(d_tile_state, d_atomic_counter),
|
|
/* start_tile, unused */ 0,
|
|
::cuda::std::move(scan_op),
|
|
init_value,
|
|
num_items,
|
|
num_stages)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Check for failure to launch
|
|
if (const auto error = CubDebug(cudaPeekAtLastError()))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
// Sync the stream if specified to flush runtime errors
|
|
if (const auto error = CubDebug(detail::DebugSyncStream(stream)))
|
|
{
|
|
return error;
|
|
}
|
|
}
|
|
#else // __cccl_ptx_isa >= 860
|
|
static_assert(sizeof(policy_getter) == 0,
|
|
"Implementation bug: Tuning policy selected lookahead, but supported PTX ISA is too low");
|
|
#endif // __cccl_ptx_isa >= 860
|
|
return cudaSuccess;
|
|
}
|
|
|
|
template <typename PolicyGetter,
|
|
typename InputIteratorT,
|
|
typename OutputIteratorT,
|
|
typename ScanOpT,
|
|
typename InitValueT,
|
|
typename OffsetT,
|
|
typename KernelSource,
|
|
typename KernelLauncherFactory>
|
|
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t invoke(
|
|
PolicyGetter policy_getter,
|
|
void* d_temp_storage,
|
|
size_t& temp_storage_bytes,
|
|
InputIteratorT d_in,
|
|
OutputIteratorT d_out,
|
|
ScanOpT scan_op,
|
|
InitValueT init_value,
|
|
OffsetT num_items,
|
|
cudaStream_t stream,
|
|
::cuda::compute_capability cc,
|
|
KernelSource kernel_source,
|
|
KernelLauncherFactory launcher_factory)
|
|
{
|
|
const bool dependent_launch = cc >= ::cuda::compute_capability{9, 0};
|
|
if CUB_DETAIL_CONSTEXPR_ISH (policy_getter().algorithm == ScanAlgorithm::lookahead)
|
|
{
|
|
const bool atomic_scheduling = cc == ::cuda::compute_capability{9, 0};
|
|
return invoke_lookahead(
|
|
policy_getter,
|
|
d_temp_storage,
|
|
temp_storage_bytes,
|
|
d_in,
|
|
d_out,
|
|
scan_op,
|
|
init_value,
|
|
num_items,
|
|
stream,
|
|
dependent_launch,
|
|
atomic_scheduling,
|
|
kernel_source,
|
|
launcher_factory);
|
|
}
|
|
else
|
|
{
|
|
return invoke_lookback(
|
|
policy_getter,
|
|
d_temp_storage,
|
|
temp_storage_bytes,
|
|
d_in,
|
|
d_out,
|
|
scan_op,
|
|
init_value,
|
|
num_items,
|
|
stream,
|
|
dependent_launch,
|
|
kernel_source,
|
|
launcher_factory);
|
|
}
|
|
}
|
|
|
|
template <
|
|
ForceInclusive EnforceInclusive = ForceInclusive::No,
|
|
bool StableReductionOrder = false,
|
|
typename InputIteratorT,
|
|
typename OutputIteratorT,
|
|
typename ScanOpT,
|
|
typename InitValueT,
|
|
typename OffsetT,
|
|
typename AccumT = ::cuda::std::__accumulator_t<ScanOpT,
|
|
cub::detail::it_value_t<InputIteratorT>,
|
|
::cuda::std::_If<::cuda::std::is_same_v<InitValueT, NullType>,
|
|
cub::detail::it_value_t<InputIteratorT>,
|
|
typename InitValueT::value_type>>,
|
|
typename PolicySelector =
|
|
policy_selector_from_types<InputIteratorT, OutputIteratorT, AccumT, OffsetT, ScanOpT, StableReductionOrder>,
|
|
typename KernelSource = DeviceScanKernelSource<
|
|
PolicySelector,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t<InputIteratorT>,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t<OutputIteratorT>,
|
|
ScanOpT,
|
|
InitValueT,
|
|
OffsetT,
|
|
AccumT,
|
|
EnforceInclusive,
|
|
StableReductionOrder>,
|
|
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
|
|
#if _CCCL_HAS_CONCEPTS()
|
|
requires scan_policy_selector<PolicySelector>
|
|
#endif // _CCCL_HAS_CONCEPTS()
|
|
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch(
|
|
void* d_temp_storage,
|
|
size_t& temp_storage_bytes,
|
|
InputIteratorT d_in,
|
|
OutputIteratorT d_out,
|
|
ScanOpT scan_op,
|
|
InitValueT init_value,
|
|
OffsetT num_items,
|
|
cudaStream_t stream,
|
|
PolicySelector policy_selector = {},
|
|
KernelSource kernel_source = {},
|
|
KernelLauncherFactory launcher_factory = {}) -> cudaError_t
|
|
{
|
|
static_assert(::cuda::std::is_unsigned_v<OffsetT> && sizeof(OffsetT) >= 4,
|
|
"DispatchScan only supports unsigned offset types of at least 4-bytes");
|
|
|
|
::cuda::compute_capability cc{};
|
|
if (const auto error = CubDebug(launcher_factory.PtxComputeCap(cc)))
|
|
{
|
|
return error;
|
|
}
|
|
|
|
#if _CCCL_HOSTED() && defined(CUB_DEBUG_LOG)
|
|
NV_IF_TARGET(NV_IS_HOST, ({
|
|
std::stringstream ss;
|
|
ss << policy_selector(cc);
|
|
_CubLog("Dispatching DeviceScan to compute capability %d.%d with tuning: %s\n",
|
|
cc.major_cap(),
|
|
cc.minor_cap(),
|
|
ss.str().c_str());
|
|
}))
|
|
#endif // _CCCL_HOSTED() && defined(CUB_DEBUG_LOG)
|
|
|
|
return dispatch_compute_cap(policy_selector, cc, [&](auto policy_getter) {
|
|
return invoke(
|
|
policy_getter,
|
|
d_temp_storage,
|
|
temp_storage_bytes,
|
|
d_in,
|
|
d_out,
|
|
scan_op,
|
|
init_value,
|
|
num_items,
|
|
stream,
|
|
cc,
|
|
kernel_source,
|
|
launcher_factory);
|
|
});
|
|
}
|
|
|
|
template <typename AccumT,
|
|
ForceInclusive EnforceInclusive = ForceInclusive::No,
|
|
bool StableReductionOrder = false,
|
|
typename InputIteratorT,
|
|
typename OutputIteratorT,
|
|
typename ScanOpT,
|
|
typename InitValueT,
|
|
typename OffsetT,
|
|
typename PolicySelector =
|
|
policy_selector_from_types<InputIteratorT, OutputIteratorT, AccumT, OffsetT, ScanOpT, StableReductionOrder>,
|
|
typename KernelSource = DeviceScanKernelSource<
|
|
PolicySelector,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t<InputIteratorT>,
|
|
THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator_t<OutputIteratorT>,
|
|
ScanOpT,
|
|
InitValueT,
|
|
OffsetT,
|
|
AccumT,
|
|
EnforceInclusive,
|
|
StableReductionOrder>,
|
|
typename KernelLauncherFactory = CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY>
|
|
CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE auto dispatch_with_accum(
|
|
void* d_temp_storage,
|
|
size_t& temp_storage_bytes,
|
|
InputIteratorT d_in,
|
|
OutputIteratorT d_out,
|
|
ScanOpT scan_op,
|
|
InitValueT init_value,
|
|
OffsetT num_items,
|
|
cudaStream_t stream,
|
|
PolicySelector policy_selector = {},
|
|
KernelSource kernel_source = {},
|
|
KernelLauncherFactory launcher_factory = {}) -> cudaError_t
|
|
{
|
|
return dispatch<EnforceInclusive,
|
|
StableReductionOrder,
|
|
InputIteratorT,
|
|
OutputIteratorT,
|
|
ScanOpT,
|
|
InitValueT,
|
|
OffsetT,
|
|
AccumT>(
|
|
d_temp_storage,
|
|
temp_storage_bytes,
|
|
d_in,
|
|
d_out,
|
|
scan_op,
|
|
init_value,
|
|
num_items,
|
|
stream,
|
|
policy_selector,
|
|
kernel_source,
|
|
launcher_factory);
|
|
}
|
|
} // namespace detail::scan
|
|
|
|
CUB_NAMESPACE_END
|