Skip to content

Integrate Reduce::ByKey to prefetching #9698

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
gonidelis wants to merge 27 commits into from
Closed

Integrate Reduce::ByKey to prefetching #9698

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
7b4b610
Add cuda::ptx::prefetch::l1/l2 wrappers for prefetching in libcudacxx
gonidelis Jun 27, 2026
25c153e
Add BLOCK_LOAD_DIRECT_PREFETCH implementation
gonidelis Jun 16, 2026
0494059
integrate reduce_by_key to use load prefetch
gonidelis Jun 16, 2026
b3ade6a
hack all workloads to use prefetch
gonidelis Jun 23, 2026
1f58325
Make BlockLoadPrefetch a compile-time parameter to BlockLoad
gonidelis Jun 24, 2026
f6fd431
enable prefetching in reduce_by_key benchmark
gonidelis Jun 24, 2026
4d41b1a
Fix ::cub:: qualified prefetch call in wrapped-namespace builds
gonidelis Jun 25, 2026
76fb6ab
expose stride as a parameter for prefetching
gonidelis Jun 26, 2026
2f64d28
Move load_prefetch after load_modifier in ReduceByKeyPolicy and other…
gonidelis Jun 26, 2026
29335c5
move blockloadprefetch in detail namespace
gonidelis Jun 26, 2026
7ecaf9b
revert load_prefetch field ordering in ReduceByKeyPolicy to not break…
gonidelis Jun 27, 2026
965b5f3
fix RleEncodePolicy adapter: drop missing load_prefetch field
gonidelis Jun 27, 2026
09270c9
remove libcudacxx ptx prefetch wrappers
gonidelis Jun 30, 2026
9e42ac4
- remove cacheloadmodified iterator handling and document it in note
gonidelis Jun 30, 2026
2664eca
reviews and ci fix
gonidelis Jun 30, 2026
82cf82b
Introduce TMA prefetch path for sm90+
gonidelis Jun 30, 2026
85dbadf
move prefetch enum place so no forward declaration is needed
gonidelis Jun 30, 2026
b97fd44
encode reducebykey prefetch (non-tma) tunings run
gonidelis Jun 30, 2026
0030373
port reducebykey prefetch tunings to the live policy_selector (new tu…
gonidelis Jun 30, 2026
c8b9cd0
use block_elect_one instead of tid=0
gonidelis Jul 1, 2026
d99482b
properly align base and size to 16bytes per ISA instructions
gonidelis Jul 1, 2026
23a17d4
enable prefetch for all BlockLoad algorithms
gonidelis Jul 2, 2026
8df100a
strip reducebyekey integration
gonidelis Jul 2, 2026
94962d3
Restore catch2_test_device_reduce_env.cu to merge-base
gonidelis Jul 2, 2026
6b12878
add tests
gonidelis Jul 2, 2026
8e46ea3
Reintroduce reduce_by_key prefetch integration on top of generalized …
gonidelis Jul 2, 2026
62def43
strip tunings from old tuning policy API
gonidelis Jul 3, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions cub/benchmarks/bench/reduce/by_key.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,11 +8,12 @@

// %RANGE% TUNE_ITEMS ipt 7:24:1
// %RANGE% TUNE_THREADS tpb 128:1024:32
// %RANGE% TUNE_TRANSPOSE trp 0:1:1
// %RANGE% TUNE_LOAD_ALGORITHM alg 0:5:1
// %RANGE% TUNE_LOAD ld 0:1:1
// %RANGE% TUNE_MAGIC_NS ns 0:2048:4
// %RANGE% TUNE_DELAY_CONSTRUCTOR_ID dcid 0:7:1
// %RANGE% TUNE_L2_WRITE_LATENCY_NS l2w 0:1200:5
// %RANGE% TUNE_LOAD_PREFETCH prf 0:3:1

#if !TUNE_BASE
struct bench_reduce_by_key_policy_selector
Expand All @@ -22,10 +23,11 @@ struct bench_reduce_by_key_policy_selector
return {
TUNE_THREADS,
TUNE_ITEMS,
TUNE_TRANSPOSE == 0 ? cub::BLOCK_LOAD_DIRECT : cub::BLOCK_LOAD_WARP_TRANSPOSE,
static_cast<cub::BlockLoadAlgorithm>(TUNE_LOAD_ALGORITHM),
TUNE_LOAD == 0 ? cub::LOAD_DEFAULT : cub::LOAD_CA,
cub::BLOCK_SCAN_WARP_SCANS,
lookback_delay_policy,
static_cast<cub::detail::BlockLoadPrefetch>(TUNE_LOAD_PREFETCH),
};
}
};
Expand Down
24 changes: 18 additions & 6 deletions cub/cub/agent/agent_reduce_by_key.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,14 +43,16 @@ template <int ThreadsPerBlock,
BlockLoadAlgorithm LoadAlgorithm,
CacheLoadModifier LoadModifier,
BlockScanAlgorithm ScanAlgorithm,
typename DelayConstructorT = detail::fixed_delay_constructor_t<350, 450>>
typename DelayConstructorT = detail::fixed_delay_constructor_t<350, 450>,
detail::BlockLoadPrefetch LoadPrefetch = detail::BlockLoadPrefetch::none>
struct agent_reduce_by_key_policy
{
static constexpr int BLOCK_THREADS = ThreadsPerBlock;
static constexpr int ITEMS_PER_THREAD = ItemsPerThread;
static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm;
static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier;
static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm;
static constexpr int BLOCK_THREADS = ThreadsPerBlock;
static constexpr int ITEMS_PER_THREAD = ItemsPerThread;
static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm;
static constexpr CacheLoadModifier LOAD_MODIFIER = LoadModifier;
static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm;
static constexpr detail::BlockLoadPrefetch LOAD_PREFETCH = LoadPrefetch;

struct detail
{
Expand Down Expand Up @@ -740,6 +742,16 @@ struct AgentReduceByKey
// Remaining items (including this tile)
OffsetT num_remaining = num_items - tile_offset;

// Fire prefetch before look-back stall so data arrives in L2 in time
if constexpr (AgentReduceByKeyPolicyT::LOAD_PREFETCH != detail::BlockLoadPrefetch::none)
{
const int items_to_prefetch = static_cast<int>(::cuda::std::min(num_remaining, static_cast<OffsetT>(TILE_ITEMS)));
detail::prefetch_block_load_tile<BLOCK_THREADS, 128, AgentReduceByKeyPolicyT::LOAD_PREFETCH>(
threadIdx.x, d_keys_in + tile_offset, items_to_prefetch);
detail::prefetch_block_load_tile<BLOCK_THREADS, 128, AgentReduceByKeyPolicyT::LOAD_PREFETCH>(
threadIdx.x, d_values_in + tile_offset, items_to_prefetch);
}
Comment on lines +745 to +753

@coderabbitai coderabbitai Bot Jul 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
# Confirm CacheModifiedInputIterator's iterator_concept / to_address support
rg -n "class CacheModifiedInputIterator" -A 80 cub/cub/iterator/cache_modified_input_iterator.cuh | rg -n "iterator_concept|iterator_category|to_address|contiguous"

Repository: NVIDIA/cccl

Length of output: 260

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Map the relevant files and read the smallest useful slices.
for f in \
  cub/cub/iterator/cache_modified_input_iterator.cuh \
  cub/cub/block/block_load.cuh \
  cub/cub/agent/agent_reduce_by_key.cuh \
  cub/cub/agent/agent_reduce_by_key.tuning.cuh \
  cub/cub/iterator/util_type.cuh
do
  echo "=== $f ==="
  wc -l "$f"
done

echo
echo "=== cache_modified_input_iterator.cuh excerpt ==="
sed -n '1,220p' cub/cub/iterator/cache_modified_input_iterator.cuh

echo
echo "=== block_load.cuh prefetch-related excerpt ==="
rg -n "prefetch|can_prefetch|contiguous_iterator|CacheModifiedInputIterator" -A 8 -B 8 cub/cub/block/block_load.cuh

echo
echo "=== agent_reduce_by_key.cuh wrapping/prefetch excerpt ==="
sed -n '180,220p' cub/cub/agent/agent_reduce_by_key.cuh
echo
sed -n '730,760p' cub/cub/agent/agent_reduce_by_key.cuh

echo
echo "=== tuning references ==="
rg -n "load_prefetch|BlockLoadPrefetch::l2|LOAD_MODIFIER|LOAD_DEFAULT" -A 4 -B 4 cub/cub/agent cub/cub | head -n 200

Repository: NVIDIA/cccl

Length of output: 551

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "=== locate tuning file ==="
fd -a 'agent_reduce_by_key' cub/cub/agent

echo
echo "=== cache_modified_input_iterator.cuh relevant lines ==="
rg -n "class CacheModifiedInputIterator|iterator_category|iterator_concept|to_address|contiguous_iterator|operator\\+|operator\\*" -A 6 -B 6 cub/cub/iterator/cache_modified_input_iterator.cuh

echo
echo "=== block_load.cuh prefetch and iterator constraints ==="
rg -n "prefetch|can_prefetch_block_load|contiguous_iterator|CacheModifiedInputIterator|load_modifier" -A 8 -B 8 cub/cub/block/block_load.cuh

echo
echo "=== agent_reduce_by_key.cuh wrapping and prefetch sites ==="
sed -n '185,215p' cub/cub/agent/agent_reduce_by_key.cuh
echo
sed -n '740,758p' cub/cub/agent/agent_reduce_by_key.cuh

Repository: NVIDIA/cccl

Length of output: 21012

important: prefetch needs the unwrapped input iterator prefetch_block_load_tile is suppressed for CacheModifiedInputIterator, so calling it on d_keys_in + tile_offset / d_values_in + tile_offset is a no-op for raw-pointer inputs. That leaves LOAD_PREFETCH ineffective on the common pointer-based DeviceReduce::ReduceByKey path; prefetch the raw pointer before wrapping, or use a separate unwrapped prefetch path.


if (num_remaining > TILE_ITEMS)
{
// Not last tile
Expand Down
204 changes: 199 additions & 5 deletions cub/cub/block/block_load.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,15 @@
#include <cub/util_ptx.cuh>
#include <cub/util_type.cuh>

#include <cuda/__cmath/round_up.h>
#include <cuda/__memcpy_async/elect_one.h>
#include <cuda/__memory/align_down.h>
#include <cuda/__ptx/ptx_helper_functions.h>
#include <cuda/std/__concepts/same_as.h>
#include <cuda/std/__fwd/format.h>
#include <cuda/std/__host_stdlib/ostream>
#include <cuda/std/__new/device_new.h>
#include <cuda/std/cstddef>

CUB_NAMESPACE_BEGIN

Expand Down Expand Up @@ -230,6 +235,120 @@ InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (&

#endif // _CCCL_DOXYGEN_INVOKED

namespace detail
{
//! @rst
//! Enumerates the cache levels that :cpp:class:`cub::BlockLoad` can prefetch into when using
//! :cpp:enumerator:`cub::BLOCK_LOAD_DIRECT`.
//!
//! Pass as the ``Prefetch`` template argument of :cpp:class:`cub::BlockLoad`. The default is
//! ``detail::BlockLoadPrefetch::none`` (no prefetch).
//!
//! .. note::
//! Prefetch hints are silently suppressed for :cpp:class:`cub::CacheModifiedInputIterator`.
//! Those iterators apply a load modifier (e.g. ``LOAD_NC``) that routes loads through a
//! specific cache path. Emitting a prefetch hint for the same addresses would touch caches
//! the user explicitly chose to bypass, defeating the purpose of the modifier.
//! Prefetch hints are also suppressed for any iterator type that is not a contiguous
//! memory iterator (i.e., one for which ``std::to_address`` is not applicable).
//!
//! @endrst
enum class BlockLoadPrefetch : int
{
//! No prefetch hint emitted. Default behavior, identical to not specifying a prefetch level.
none,
//! Emit an L2 prefetch hint (``prefetch.global.L2``) before loading each cache line.
//! All threads collectively cover the tile, each issuing one instruction per cache line.
l2,
//! Emit an L1 prefetch hint before loading each cache line. Falls back to L2 on architectures
//! that do not support real L1 prefetch.
l1,
//! Emit a TMA bulk prefetch into L2 (``cp.async.bulk.prefetch``) before loading the tile.
//! One elected thread issues a single instruction for the whole tile via the TMA engine,
//! leaving the SM load/store units free. Requires SM_90 or later (Hopper/Blackwell).
//! Falls back to a no-op on older architectures.
bulk_l2,
};

#if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED)
inline ::std::ostream& operator<<(::std::ostream& os, BlockLoadPrefetch prefetch)
{
switch (prefetch)
{
case BlockLoadPrefetch::none:
return os << "detail::BlockLoadPrefetch::none";
case BlockLoadPrefetch::l2:
return os << "detail::BlockLoadPrefetch::l2";
case BlockLoadPrefetch::l1:
return os << "detail::BlockLoadPrefetch::l1";
case BlockLoadPrefetch::bulk_l2:
return os << "detail::BlockLoadPrefetch::bulk_l2";
default:
return os << "<unknown detail::BlockLoadPrefetch: " << static_cast<int>(prefetch) << ">";
}
}
#endif // _CCCL_HOSTED() && !_CCCL_DOXYGEN_INVOKED

template <typename RandomAccessIterator>
inline constexpr bool can_prefetch_block_load =
::cuda::std::contiguous_iterator<RandomAccessIterator> && ::cuda::std::__can_to_address<RandomAccessIterator>;

template <int ThreadsPerBlock,
int PrefetchStride = 128,
BlockLoadPrefetch Prefetch = BlockLoadPrefetch::l2,
typename RandomAccessIterator>
_CCCL_DEVICE _CCCL_FORCEINLINE void
prefetch_block_load_tile(int linear_tid, RandomAccessIterator block_src_it, int items_to_prefetch)
{
if constexpr (can_prefetch_block_load<RandomAccessIterator>)
{
using input_t = cub::detail::it_value_t<RandomAccessIterator>;
const unsigned int total_bytes = static_cast<unsigned int>(items_to_prefetch) * sizeof(input_t);
const auto* const base = reinterpret_cast<const char*>(::cuda::std::to_address(block_src_it));

if constexpr (Prefetch == BlockLoadPrefetch::bulk_l2)
{
// One elected thread issues a single TMA bulk prefetch for the whole tile.
// cp.async.bulk.prefetch is fire-and-forget: no commit_group/wait_group needed.
// Requires SM_90+; on older archs __block_elect_one() falls back to threadIdx.x == 0.
if (::cuda::device::__block_elect_one())
{
// srcMem must be 16-byte aligned per PTX ISA; align base down and extend size to compensate
const auto* const aligned_base = ::cuda::align_down(base, 16);
const unsigned int prefix = static_cast<unsigned int>(base - aligned_base);
const unsigned int aligned_size = ::cuda::round_up(total_bytes + prefix, 16u);
if (aligned_size > 0)
{
#if _CCCL_CUDA_COMPILER(NVHPC) || __CUDA_ARCH__ >= 900
asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;"
:
: "l"(::cuda::ptx::__as_ptr_gmem(aligned_base)), "r"(aligned_size)
: "memory");
#endif
}
}
}
else
{
_CCCL_PRAGMA_NOUNROLL()
for (unsigned int offset = static_cast<unsigned int>(linear_tid) * PrefetchStride; offset < total_bytes;
offset += static_cast<unsigned int>(ThreadsPerBlock) * PrefetchStride)
{
// TODO: replace with cuda::ptx::prefetch_L1/L2 once exposed in libcudacxx
if constexpr (Prefetch == BlockLoadPrefetch::l1)
{
asm volatile("prefetch.global.L1 [%0];" : : "l"(::cuda::ptx::__as_ptr_gmem(base + offset)) : "memory");
}
else
{
asm volatile("prefetch.global.L2 [%0];" : : "l"(::cuda::ptx::__as_ptr_gmem(base + offset)) : "memory");
}
}
}
}
}
} // namespace detail

//! @rst
//! Load a linear segment of items into a blocked arrangement across the thread block.
//!
Expand Down Expand Up @@ -788,6 +907,7 @@ CUB_NAMESPACE_BEGIN
//!
//! #. :cpp:enumerator:`cub::BLOCK_LOAD_DIRECT`:
//! A :ref:`blocked arrangement <flexible-data-arrangement>` of data is read directly from memory.
//! Combine with ``Prefetch = cub::detail::BlockLoadPrefetch::l2`` to emit L2 prefetch hints before each load.
//! #. :cpp:enumerator:`cub::BLOCK_LOAD_STRIPED`:
//! A :ref:`striped arrangement <flexible-data-arrangement>` of data is read directly from memory.
//! #. :cpp:enumerator:`cub::BLOCK_LOAD_VECTORIZE`:
Expand Down Expand Up @@ -865,9 +985,10 @@ CUB_NAMESPACE_BEGIN
template <typename T,
int BlockDimX,
int ItemsPerThread,
BlockLoadAlgorithm Algorithm = BLOCK_LOAD_DIRECT,
int BlockDimY = 1,
int BlockDimZ = 1>
BlockLoadAlgorithm Algorithm = BLOCK_LOAD_DIRECT,
int BlockDimY = 1,
int BlockDimZ = 1,
detail::BlockLoadPrefetch Prefetch = detail::BlockLoadPrefetch::none>
class BlockLoad
{
static constexpr int ThreadsPerBlock = BlockDimX * BlockDimY * BlockDimZ; // total threads in the block
Expand Down Expand Up @@ -995,14 +1116,29 @@ public:
{
if constexpr (Algorithm == BLOCK_LOAD_DIRECT)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(
linear_tid, block_src_it, ThreadsPerBlock * ItemsPerThread);
}
LoadDirectBlocked(linear_tid, block_src_it, dst_items);
}
else if constexpr (Algorithm == BLOCK_LOAD_STRIPED)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(
linear_tid, block_src_it, ThreadsPerBlock * ItemsPerThread);
}
LoadDirectStriped<ThreadsPerBlock>(linear_tid, block_src_it, dst_items);
}
else if constexpr (Algorithm == BLOCK_LOAD_VECTORIZE)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(
linear_tid, block_src_it, ThreadsPerBlock * ItemsPerThread);
}
if constexpr (detail::is_CacheModifiedInputIterator<RandomAccessIterator>)
{
InternalLoadDirectBlockedVectorized<RandomAccessIterator::__modifier>(linear_tid, block_src_it.ptr, dst_items);
Expand All @@ -1019,11 +1155,21 @@ public:
}
else if constexpr (Algorithm == BLOCK_LOAD_TRANSPOSE)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(
linear_tid, block_src_it, ThreadsPerBlock * ItemsPerThread);
}
LoadDirectStriped<ThreadsPerBlock>(linear_tid, block_src_it, dst_items);
block_exchange(temp_storage).StripedToBlocked(dst_items, dst_items);
}
else if constexpr (Algorithm == BLOCK_LOAD_WARP_TRANSPOSE || Algorithm == BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(
linear_tid, block_src_it, ThreadsPerBlock * ItemsPerThread);
}
LoadDirectWarpStriped(linear_tid, block_src_it, dst_items);
block_exchange(temp_storage).WarpStripedToBlocked(dst_items, dst_items);
}
Expand Down Expand Up @@ -1082,21 +1228,45 @@ public:
_CCCL_DEVICE _CCCL_FORCEINLINE void
Load(RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread], int block_items_end)
{
if constexpr (Algorithm == BLOCK_LOAD_DIRECT || Algorithm == BLOCK_LOAD_VECTORIZE)
if constexpr (Algorithm == BLOCK_LOAD_DIRECT)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end);
}
else if constexpr (Algorithm == BLOCK_LOAD_VECTORIZE)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end);
}
else if constexpr (Algorithm == BLOCK_LOAD_STRIPED)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectStriped<ThreadsPerBlock>(linear_tid, block_src_it, dst_items, block_items_end);
}
else if constexpr (Algorithm == BLOCK_LOAD_TRANSPOSE)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectStriped<ThreadsPerBlock>(linear_tid, block_src_it, dst_items, block_items_end);
block_exchange(temp_storage).StripedToBlocked(dst_items, dst_items);
}
else if constexpr (Algorithm == BLOCK_LOAD_WARP_TRANSPOSE || Algorithm == BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectWarpStriped(linear_tid, block_src_it, dst_items, block_items_end);
block_exchange(temp_storage).WarpStripedToBlocked(dst_items, dst_items);
}
Expand Down Expand Up @@ -1158,21 +1328,45 @@ public:
_CCCL_DEVICE _CCCL_FORCEINLINE void
Load(RandomAccessIterator block_src_it, T (&dst_items)[ItemsPerThread], int block_items_end, DefaultT oob_default)
{
if constexpr (Algorithm == BLOCK_LOAD_DIRECT || Algorithm == BLOCK_LOAD_VECTORIZE)
if constexpr (Algorithm == BLOCK_LOAD_DIRECT)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end, oob_default);
}
else if constexpr (Algorithm == BLOCK_LOAD_VECTORIZE)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectBlocked(linear_tid, block_src_it, dst_items, block_items_end, oob_default);
}
else if constexpr (Algorithm == BLOCK_LOAD_STRIPED)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectStriped<ThreadsPerBlock>(linear_tid, block_src_it, dst_items, block_items_end, oob_default);
}
else if constexpr (Algorithm == BLOCK_LOAD_TRANSPOSE)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectStriped<ThreadsPerBlock>(linear_tid, block_src_it, dst_items, block_items_end, oob_default);
block_exchange(temp_storage).StripedToBlocked(dst_items, dst_items);
}
else if constexpr (Algorithm == BLOCK_LOAD_WARP_TRANSPOSE || Algorithm == BLOCK_LOAD_WARP_TRANSPOSE_TIMESLICED)
{
if constexpr (Prefetch != detail::BlockLoadPrefetch::none)
{
detail::prefetch_block_load_tile<ThreadsPerBlock, 128, Prefetch>(linear_tid, block_src_it, block_items_end);
}
LoadDirectWarpStriped(linear_tid, block_src_it, dst_items, block_items_end, oob_default);
block_exchange(temp_storage).WarpStripedToBlocked(dst_items, dst_items);
}
Expand Down
6 changes: 4 additions & 2 deletions cub/cub/device/dispatch/dispatch_reduce_by_key.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,8 @@ __launch_bounds__(int(current_policy<PolicySelector>().threads_per_block))
policy.load_algorithm,
policy.load_modifier,
policy.scan_algorithm,
delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>>;
delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>,
policy.load_prefetch>;

using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<
AgentReduceByKeyPolicyT,
Expand Down Expand Up @@ -654,7 +655,8 @@ _CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_get
policy.load_algorithm,
policy.load_modifier,
policy.scan_algorithm,
delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>>;
delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>,
policy.load_prefetch>;
using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<Policy, AgentReduceByKey, Args...>;
return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS,
vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD,
Expand Down
Loading
Loading