cub/block/block_scan.cuh
File members: cub/block/block_scan.cuh
/******************************************************************************
* Copyright (c) 2011, Duane Merrill. All rights reserved.
* Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
* * Redistributions of source code must retain the above copyright
* notice, this list of conditions and the following disclaimer.
* * Redistributions in binary form must reproduce the above copyright
* notice, this list of conditions and the following disclaimer in the
* documentation and/or other materials provided with the distribution.
* * Neither the name of the NVIDIA CORPORATION nor the
* names of its contributors may be used to endorse or promote products
* derived from this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
* ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
* WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
* DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
* (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
* LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
* ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
* (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
* SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
******************************************************************************/
#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_scan_raking.cuh>
#include <cub/block/specializations/block_scan_warp_scans.cuh>
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>
#include <cuda/std/type_traits>
CUB_NAMESPACE_BEGIN
/******************************************************************************
* Algorithmic variants
******************************************************************************/
enum BlockScanAlgorithm
{
BLOCK_SCAN_RAKING,
BLOCK_SCAN_RAKING_MEMOIZE,
BLOCK_SCAN_WARP_SCANS,
};
template <typename T,
int BLOCK_DIM_X,
BlockScanAlgorithm ALGORITHM = BLOCK_SCAN_RAKING,
int BLOCK_DIM_Y = 1,
int BLOCK_DIM_Z = 1,
int LEGACY_PTX_ARCH = 0>
class BlockScan
{
private:
enum
{
BLOCK_THREADS = BLOCK_DIM_X * BLOCK_DIM_Y * BLOCK_DIM_Z,
};
static constexpr BlockScanAlgorithm SAFE_ALGORITHM =
((ALGORITHM == BLOCK_SCAN_WARP_SCANS) && (BLOCK_THREADS % CUB_WARP_THREADS(0) != 0))
? BLOCK_SCAN_RAKING
: ALGORITHM;
using WarpScans = BlockScanWarpScans<T, BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z>;
using Raking =
BlockScanRaking<T, BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z, (SAFE_ALGORITHM == BLOCK_SCAN_RAKING_MEMOIZE)>;
using InternalBlockScan = ::cuda::std::_If<SAFE_ALGORITHM == BLOCK_SCAN_WARP_SCANS, WarpScans, Raking>;
using _TempStorage = typename InternalBlockScan::TempStorage;
_TempStorage& temp_storage;
unsigned int linear_tid;
_CCCL_DEVICE _CCCL_FORCEINLINE _TempStorage& PrivateStorage()
{
__shared__ _TempStorage private_storage;
return private_storage;
}
public:
struct TempStorage : Uninitialized<_TempStorage>
{};
_CCCL_DEVICE _CCCL_FORCEINLINE BlockScan()
: temp_storage(PrivateStorage())
, linear_tid(RowMajorTid(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z))
{}
_CCCL_DEVICE _CCCL_FORCEINLINE BlockScan(TempStorage& temp_storage)
: temp_storage(temp_storage.Alias())
, linear_tid(RowMajorTid(BLOCK_DIM_X, BLOCK_DIM_Y, BLOCK_DIM_Z))
{}
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output)
{
T initial_value{};
ExclusiveScan(input, output, initial_value, ::cuda::std::plus<>{});
}
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output, T& block_aggregate)
{
T initial_value{};
ExclusiveScan(input, output, initial_value, ::cuda::std::plus<>{}, block_aggregate);
}
template <typename BlockPrefixCallbackOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T input, T& output, BlockPrefixCallbackOp& block_prefix_callback_op)
{
ExclusiveScan(input, output, ::cuda::std::plus<>{}, block_prefix_callback_op);
}
template <int ITEMS_PER_THREAD>
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD])
{
T initial_value{};
ExclusiveScan(input, output, initial_value, ::cuda::std::plus<>{});
}
template <int ITEMS_PER_THREAD>
_CCCL_DEVICE _CCCL_FORCEINLINE void
ExclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T& block_aggregate)
{
// Reduce consecutive thread items in registers
T initial_value{};
ExclusiveScan(input, output, initial_value, ::cuda::std::plus<>{}, block_aggregate);
}
template <int ITEMS_PER_THREAD, typename BlockPrefixCallbackOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveSum(
T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], BlockPrefixCallbackOp& block_prefix_callback_op)
{
ExclusiveScan(input, output, ::cuda::std::plus<>{}, block_prefix_callback_op);
}
template <typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, T initial_value, ScanOp scan_op)
{
InternalBlockScan(temp_storage).ExclusiveScan(input, output, initial_value, scan_op);
}
template <typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
ExclusiveScan(T input, T& output, T initial_value, ScanOp scan_op, T& block_aggregate)
{
InternalBlockScan(temp_storage).ExclusiveScan(input, output, initial_value, scan_op, block_aggregate);
}
template <typename ScanOp, typename BlockPrefixCallbackOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
ExclusiveScan(T input, T& output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op)
{
InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op, block_prefix_callback_op);
}
template <int ITEMS_PER_THREAD, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op)
{
// Reduce consecutive thread items in registers
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op);
// Exclusive scan in registers with prefix as seed
internal::ThreadScanExclusive(input, output, scan_op, thread_prefix);
}
template <int ITEMS_PER_THREAD, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(
T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op, T& block_aggregate)
{
// Reduce consecutive thread items in registers
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op, block_aggregate);
// Exclusive scan in registers with prefix as seed
internal::ThreadScanExclusive(input, output, scan_op, thread_prefix);
}
template <int ITEMS_PER_THREAD, typename ScanOp, typename BlockPrefixCallbackOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(
T (&input)[ITEMS_PER_THREAD],
T (&output)[ITEMS_PER_THREAD],
ScanOp scan_op,
BlockPrefixCallbackOp& block_prefix_callback_op)
{
// Reduce consecutive thread items in registers
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_prefix_callback_op);
// Exclusive scan in registers with prefix as seed
internal::ThreadScanExclusive(input, output, scan_op, thread_prefix);
}
#ifndef _CCCL_DOXYGEN_INVOKED // Do not document no-initial-value scans
template <typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, ScanOp scan_op)
{
InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op);
}
template <typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void ExclusiveScan(T input, T& output, ScanOp scan_op, T& block_aggregate)
{
InternalBlockScan(temp_storage).ExclusiveScan(input, output, scan_op, block_aggregate);
}
template <int ITEMS_PER_THREAD, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op)
{
// Reduce consecutive thread items in registers
T thread_partial = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_partial, thread_partial, scan_op);
// Exclusive scan in registers with prefix
internal::ThreadScanExclusive(input, output, scan_op, thread_partial, (linear_tid != 0));
}
template <int ITEMS_PER_THREAD, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
ExclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op, T& block_aggregate)
{
// Reduce consecutive thread items in registers
T thread_partial = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_partial, thread_partial, scan_op, block_aggregate);
// Exclusive scan in registers with prefix
internal::ThreadScanExclusive(input, output, scan_op, thread_partial, (linear_tid != 0));
}
#endif // _CCCL_DOXYGEN_INVOKED // Do not document no-initial-value scans
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output)
{
InclusiveScan(input, output, ::cuda::std::plus<>{});
}
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output, T& block_aggregate)
{
InclusiveScan(input, output, ::cuda::std::plus<>{}, block_aggregate);
}
template <typename BlockPrefixCallbackOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T input, T& output, BlockPrefixCallbackOp& block_prefix_callback_op)
{
InclusiveScan(input, output, ::cuda::std::plus<>{}, block_prefix_callback_op);
}
template <int ITEMS_PER_THREAD>
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD])
{
if (ITEMS_PER_THREAD == 1)
{
InclusiveSum(input[0], output[0]);
}
else
{
// Reduce consecutive thread items in registers
::cuda::std::plus<> scan_op;
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveSum(thread_prefix, thread_prefix);
// Inclusive scan in registers with prefix as seed
internal::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0));
}
}
template <int ITEMS_PER_THREAD>
_CCCL_DEVICE _CCCL_FORCEINLINE void
InclusiveSum(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T& block_aggregate)
{
if (ITEMS_PER_THREAD == 1)
{
InclusiveSum(input[0], output[0], block_aggregate);
}
else
{
// Reduce consecutive thread items in registers
::cuda::std::plus<> scan_op;
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveSum(thread_prefix, thread_prefix, block_aggregate);
// Inclusive scan in registers with prefix as seed
internal::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0));
}
}
template <int ITEMS_PER_THREAD, typename BlockPrefixCallbackOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveSum(
T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], BlockPrefixCallbackOp& block_prefix_callback_op)
{
if (ITEMS_PER_THREAD == 1)
{
InclusiveSum(input[0], output[0], block_prefix_callback_op);
}
else
{
// Reduce consecutive thread items in registers
::cuda::std::plus<> scan_op;
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveSum(thread_prefix, thread_prefix, block_prefix_callback_op);
// Inclusive scan in registers with prefix as seed
internal::ThreadScanInclusive(input, output, scan_op, thread_prefix);
}
}
template <typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& output, ScanOp scan_op)
{
InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op);
}
template <typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(T input, T& output, ScanOp scan_op, T& block_aggregate)
{
InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op, block_aggregate);
}
template <typename ScanOp, typename BlockPrefixCallbackOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
InclusiveScan(T input, T& output, ScanOp scan_op, BlockPrefixCallbackOp& block_prefix_callback_op)
{
InternalBlockScan(temp_storage).InclusiveScan(input, output, scan_op, block_prefix_callback_op);
}
template <int ITEMS_PER_THREAD, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op)
{
if (ITEMS_PER_THREAD == 1)
{
InclusiveScan(input[0], output[0], scan_op);
}
else
{
// Reduce consecutive thread items in registers
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_prefix, thread_prefix, scan_op);
// Inclusive scan in registers with prefix as seed (first thread does not seed)
internal::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0));
}
}
template <int ITEMS_PER_THREAD, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op)
{
// Reduce consecutive thread items in registers
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op);
// Exclusive scan in registers with prefix as seed
internal::ThreadScanInclusive(input, output, scan_op, thread_prefix);
}
template <int ITEMS_PER_THREAD, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void
InclusiveScan(T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], ScanOp scan_op, T& block_aggregate)
{
if (ITEMS_PER_THREAD == 1)
{
InclusiveScan(input[0], output[0], scan_op, block_aggregate);
}
else
{
// Reduce consecutive thread items in registers
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan (with no initial value)
ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_aggregate);
// Inclusive scan in registers with prefix as seed (first thread does not seed)
internal::ThreadScanInclusive(input, output, scan_op, thread_prefix, (linear_tid != 0));
}
}
template <int ITEMS_PER_THREAD, typename ScanOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(
T (&input)[ITEMS_PER_THREAD], T (&output)[ITEMS_PER_THREAD], T initial_value, ScanOp scan_op, T& block_aggregate)
{
// Reduce consecutive thread items in registers
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_prefix, thread_prefix, initial_value, scan_op, block_aggregate);
// Exclusive scan in registers with prefix as seed
internal::ThreadScanInclusive(input, output, scan_op, thread_prefix);
}
template <int ITEMS_PER_THREAD, typename ScanOp, typename BlockPrefixCallbackOp>
_CCCL_DEVICE _CCCL_FORCEINLINE void InclusiveScan(
T (&input)[ITEMS_PER_THREAD],
T (&output)[ITEMS_PER_THREAD],
ScanOp scan_op,
BlockPrefixCallbackOp& block_prefix_callback_op)
{
if (ITEMS_PER_THREAD == 1)
{
InclusiveScan(input[0], output[0], scan_op, block_prefix_callback_op);
}
else
{
// Reduce consecutive thread items in registers
T thread_prefix = cub::ThreadReduce(input, scan_op);
// Exclusive thread block-scan
ExclusiveScan(thread_prefix, thread_prefix, scan_op, block_prefix_callback_op);
// Inclusive scan in registers with prefix as seed
internal::ThreadScanInclusive(input, output, scan_op, thread_prefix);
}
}
};
CUB_NAMESPACE_END