Skip to content

[CUB] Use the vectorized load machinery in BlockLoad and WarpLoad #12213

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

Open
miscco wants to merge 1 commit into from
+257 −68
Open

[CUB] Use the vectorized load machinery in BlockLoad and WarpLoad #12213

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
30 changes: 27 additions & 3 deletions cub/cub/agent/agent_topk.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,10 @@
#include <cub/detail/topk_helpers.cuh>
#include <cub/detail/uninitialized.cuh>

#include <thrust/type_traits/is_contiguous_iterator.h>

#include <cuda/__cmath/ceil_div.h>
#include <cuda/std/__memory/pointer_traits.h>
#include <cuda/std/__type_traits/is_same.h>

CUB_NAMESPACE_BEGIN
Expand All @@ -47,18 +50,23 @@ namespace detail::topk
//! @tparam ScanAlgorithm
//! The BlockScan algorithm to use
//!
//! @tparam VecSize
//! Number of keys per vectorized load of contiguous input with BLOCK_LOAD_VECTORIZE, 0 to load through BlockLoad
//!
template <int ThreadsPerBlock,
int ItemsPerThread,
int BitsPerPass,
BlockLoadAlgorithm LoadAlgorithm,
BlockScanAlgorithm ScanAlgorithm>
BlockScanAlgorithm ScanAlgorithm,
int VecSize>
struct agent_topk_policy
{
static constexpr int threads_per_block = ThreadsPerBlock;
static constexpr int items_per_thread = ItemsPerThread;
static constexpr int bits_per_pass = BitsPerPass;
static constexpr BlockLoadAlgorithm load_algorithm = LoadAlgorithm;
static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm;
static constexpr int vec_size = VecSize;
};

template <typename KeyInT, typename OffsetT, typename OutOffsetT>
Expand Down Expand Up @@ -258,6 +266,22 @@ struct AgentTopK
// Utility methods for device topK
//---------------------------------------------------------------------

// Load a full tile in a blocked arrangement
template <typename InputItT>
_CCCL_DEVICE _CCCL_FORCEINLINE void load_tile(InputItT block_in, key_in_t (&thread_data)[items_per_thread])
{
if constexpr (AgentTopKPolicyT::load_algorithm == BLOCK_LOAD_VECTORIZE && AgentTopKPolicyT::vec_size != 0
&& THRUST_NS_QUALIFIER::is_contiguous_iterator_v<InputItT>)
{
LoadDirectBlockedVectorized<LOAD_DEFAULT, AgentTopKPolicyT::vec_size>(
static_cast<int>(threadIdx.x), ::cuda::std::to_address(block_in), thread_data);
}
else
{
block_load_input_t(temp_storage.load_input).Load(block_in, thread_data);
}
}

// Process a range of input data in tiles, calling f(key, index) for each element
template <typename InputItT, typename FuncT>
_CCCL_DEVICE _CCCL_FORCEINLINE void process_range(InputItT in, const OffsetT num_items, FuncT f)
Expand All @@ -280,7 +304,7 @@ struct AgentTopK
// Ensure that the temporary storage from previous iteration can be reused
__syncthreads();

block_load_input_t(temp_storage.load_input).Load(in + tile_base, thread_data);
load_tile(in + tile_base, thread_data);
for (int j = 0; j < items_per_thread; ++j)
{
f(thread_data[j], offset + j);
Expand All @@ -297,7 +321,7 @@ struct AgentTopK

if (num_remaining_elements == 0)
{
block_load_input_t(temp_storage.load_input).Load(in + tile_base, thread_data);
load_tile(in + tile_base, thread_data);
}
else
{
Expand Down
165 changes: 113 additions & 52 deletions cub/cub/block/block_load.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,19 +23,21 @@
#include <cub/detail/cub_vector.cuh>
#include <cub/detail/it_traits.cuh>
#include <cub/detail/null_type.cuh>
#include <cub/detail/numeric_traits.cuh>
#include <cub/detail/uninitialized.cuh>
#include <cub/detail/unit_word.cuh>
#include <cub/iterator/cache_modified_input_iterator.cuh>
#include <cub/util_ptx.cuh>

#include <cuda/__memory/__vectorized_storage.h>
#include <cuda/__type_traits/is_trivially_copyable.h>
#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>
#include <cuda/std/__new/device_new.h>
#include <cuda/std/__type_traits/is_same.h>
#include <cuda/std/__type_traits/is_volatile.h>
#include <cuda/std/__type_traits/remove_reference.h>

#include <nv/target>

CUB_NAMESPACE_BEGIN

Expand Down Expand Up @@ -182,61 +184,51 @@ _CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectBlocked(

#ifndef _CCCL_DOXYGEN_INVOKED // Do not document

//! @brief Internal implementation for load vectorization
namespace detail
{
//! Whether the storage of ``RandomAccessIterator`` can be loaded with vectors
// TODO(bgruber): remove the check for is_primitive<StorageT> in CCCL 4.0
template <typename RandomAccessIterator, typename StorageT = it_value_t<RandomAccessIterator>>
inline constexpr bool block_load_can_vectorize_v =
(is_primitive<StorageT>::value || ::cuda::is_trivially_copyable_v<StorageT>) //
&&!::cuda::std::is_volatile_v<StorageT>
&& THRUST_NS_QUALIFIER::is_contiguous_iterator_v<RandomAccessIterator>;

//! @brief Internal implementation of LoadDirectBlockedVectorized for input aligned to the vector type
//!
//! @param[in] linear_tid
//! A suitable 1D thread-identifier for the calling thread (e.g., `(threadIdx.y * blockDim.x) + linear_tid` for 2D
//! thread blocks)
//!
//! @param[in] block_src_ptr
//! Input pointer for loading from
//! Input pointer for loading from, aligned to ``cuda::__as_vectorized_storage_t<InputT, ItemsPerThread>``
//!
//! @param[out] dst_items
//! Destination to load data into
template <CacheLoadModifier MODIFIER, typename T, int ItemsPerThread>
template <CacheLoadModifier MODIFIER,
::cuda::std::size_t MaxVectorAccess = ::cuda::__max_vector_access,
typename T,
int ItemsPerThread,
typename InputT>
_CCCL_DEVICE _CCCL_FORCEINLINE void
InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (&dst_items)[ItemsPerThread])
InternalLoadDirectBlockedVectorized(int linear_tid, const InputT* block_src_ptr, T (&dst_items)[ItemsPerThread])
{
// Find biggest memory access word that T is a whole multiple of
using device_word_t = typename UnitWord<T>::DeviceWord;
_CCCL_DIAG_PUSH
# if _CCCL_COMPILER(CLANG, >=, 10)
_CCCL_DIAG_SUPPRESS_CLANG("-Wsizeof-array-div")
# endif // _CCCL_COMPILER(CLANG, >=, 10)
// NOLINTNEXTLINE(bugprone-sizeof-expression)
constexpr int total_words = static_cast<int>(sizeof(dst_items) / sizeof(device_word_t));
_CCCL_DIAG_POP
constexpr int vector_size = (total_words % 4 == 0) ? 4 : (total_words % 2 == 0) ? 2 : 1;
constexpr int vectors_per_thread = total_words / vector_size;

// Load into an array of vectors in thread-blocked order
using vector_t = typename CubVector<device_word_t, vector_size>::Type;

// Add the alignment check to ensure the vectorized loading can proceed.
if (::cuda::std::is_sufficiently_aligned<alignof(vector_t)>(block_src_ptr))
{
vector_t vec_items[vectors_per_thread];
// Load into an array of vectors in thread-blocked order
const vector_t* vec_ptr = reinterpret_cast<const vector_t*>(block_src_ptr) + linear_tid * vectors_per_thread;
// All items of a thread are loaded with the widest vectors that evenly divide them
using vector_t = ::cuda::__as_vectorized_storage_t<InputT, ItemsPerThread, MaxVectorAccess>;
_CCCL_ASSERT(::cuda::std::is_sufficiently_aligned<alignof(vector_t)>(block_src_ptr),
"block_src_ptr must be aligned to the vector type");

_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < vectors_per_thread; i++)
{
vec_items[i] = ThreadLoad<MODIFIER>(vec_ptr + i);
}
// Load in thread-blocked order
const vector_t vec_items = ThreadLoad<MODIFIER>(reinterpret_cast<const vector_t*>(block_src_ptr) + linear_tid);
Comment thread
miscco marked this conversation as resolved.

// Copy to destination
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < ItemsPerThread; i++)
{
dst_items[i] = *(reinterpret_cast<T*>(vec_items) + i);
}
}
else
// Copy to destination
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < ItemsPerThread; i++)
{
LoadDirectBlocked(linear_tid, block_src_ptr, dst_items);
dst_items[i] = vec_items.template __alias<InputT>(i);

@coderabbitai coderabbitai Bot Oct 24, 2026 •
edited
Loading

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.

Minor · Functional Correctness · Quick win

important: Convert user-defined inputs from the source object.

When InputT differs from T, this conversion runs on vec_items, not on the input object. A trivially copyable InputT can still define an operator T() that increments a mutable field or depends on its address. The new vectorized path changes that conversion’s result or side effects; the previous dispatch used direct source-object loads. Keep direct loading for user-defined conversions. (raw.githubusercontent.com)

@miscco miscco Oct 28, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

That is not correct, we are aliasing as InputT

@coderabbitai coderabbitai Bot Oct 28, 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.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
fd '^as_bytes_vectorized\.h$' . --exec ast-grep outline {} --items all
rg -n -C 10 '__alias\s*\(' libcudacxx/include/cuda/__memory/as_bytes_vectorized.h

Length of output: 3224

@miscco important: __alias<InputT> preserves the type, but not the original object’s identity. The implementation aliases vec_items’s __storage_, not block_src_ptr.

Ordinary value conversions work as intended. My concern is limited to conversions that depend on this or modify mutable source state. Aliasing the loaded copy as InputT does not preserve those behaviors.

📝 I couldn't check this conversation for new learnings, so no new learning was saved.

You are interacting with an AI system.

}
}
} // namespace detail

#endif // _CCCL_DOXYGEN_INVOKED

Expand All @@ -258,12 +250,23 @@ InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (&
//!
//! @endrst
//!
//! @tparam Modifier
//! **[optional]** Cache load modifier for the vectorized loads (default: ``LOAD_DEFAULT``)
//!
//! @tparam VecSize
//! **[optional]** The number of consecutive items in each vector. Must evenly divide ``ItemsPerThread``. Each thread
//! loads its items with ``ItemsPerThread / VecSize`` vectors, and ``0`` loads them with a single vector
//! (default: ``0``)
//!
//! @tparam T
//! **[inferred]** The data type to load.
//!
//! @tparam ItemsPerThread
//! **[inferred]** The number of consecutive items partitioned onto each thread.
//!
//! @tparam InputT
//! **[inferred]** The value type of the input, which is loaded with vectors and converted to ``T``.
//!
//! @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)
Expand All @@ -273,12 +276,75 @@ InternalLoadDirectBlockedVectorized(int linear_tid, const T* block_src_ptr, T (&
//!
//! @param[out] dst_items
//! destination to load data into
template <CacheLoadModifier Modifier = LOAD_DEFAULT, int VecSize = 0, typename T, int ItemsPerThread, typename InputT>
_CCCL_DEVICE _CCCL_FORCEINLINE void
LoadDirectBlockedVectorized(int linear_tid, const InputT* block_src_ptr, T (&dst_items)[ItemsPerThread])
{
constexpr int vec_size = VecSize == 0 ? ItemsPerThread : VecSize;
static_assert(ItemsPerThread % vec_size == 0, "VecSize must evenly divide ItemsPerThread");

if constexpr (detail::block_load_can_vectorize_v<InputT*, InputT>)
{
// Vector i of a thread is the blocked vector of thread linear_tid * vectors_per_thread + i
constexpr int vectors_per_thread = ItemsPerThread / vec_size;

// Only sm_100 and newer load vectors wider than 16 bytes with a single access, and only without a cache modifier
using vector16_t = ::cuda::__as_vectorized_storage_t<InputT, vec_size, 16ull>;

#if __cccl_ptx_isa >= 900ULL
using vector32_t = ::cuda::__as_vectorized_storage_t<InputT, vec_size, 32ull>;
if constexpr (Modifier == LOAD_DEFAULT && alignof(vector16_t) < alignof(vector32_t))
{
bool is_vector32_aligned = false;
NV_IF_TARGET(NV_PROVIDES_SM_100,
(is_vector32_aligned = ::cuda::std::is_sufficiently_aligned<alignof(vector32_t)>(block_src_ptr);))
if (is_vector32_aligned)
{
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < vectors_per_thread; i++)
{
detail::InternalLoadDirectBlockedVectorized<Modifier, 32ull>(
linear_tid * vectors_per_thread + i,
block_src_ptr,
reinterpret_cast<T(&)[vec_size]>(dst_items[i * vec_size]));
}
return;
}
}
#endif // __cccl_ptx_isa >= 900ULL

// Add the alignment check to ensure the vectorized loading can proceed.
if (::cuda::std::is_sufficiently_aligned<alignof(vector16_t)>(block_src_ptr))
{
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < vectors_per_thread; i++)
{
detail::InternalLoadDirectBlockedVectorized<Modifier, 16ull>(
linear_tid * vectors_per_thread + i, block_src_ptr, reinterpret_cast<T(&)[vec_size]>(dst_items[i * vec_size]));
}
return;
}
}

if constexpr (Modifier != LOAD_DEFAULT)
{
LoadDirectBlocked(linear_tid, CacheModifiedInputIterator<Modifier, InputT, int>(block_src_ptr), dst_items);
}
else
{
LoadDirectBlocked(linear_tid, block_src_ptr, dst_items);
}
}

#ifndef _CCCL_DOXYGEN_INVOKED // Do not document
// Accepts explicitly specified T and ItemsPerThread, which cannot bind to the leading Modifier of the overload above
template <typename T, int ItemsPerThread>
_CCCL_DEVICE _CCCL_FORCEINLINE void
LoadDirectBlockedVectorized(int linear_tid, T* block_src_ptr, T (&dst_items)[ItemsPerThread])
{
InternalLoadDirectBlockedVectorized<LOAD_DEFAULT>(linear_tid, block_src_ptr, dst_items);
LoadDirectBlockedVectorized<LOAD_DEFAULT>(linear_tid, block_src_ptr, dst_items);
}
#endif // _CCCL_DOXYGEN_INVOKED

//! @}
//! @name Striped arrangement I/O (direct)
Expand Down Expand Up @@ -1015,18 +1081,13 @@ public:
}
else if constexpr (Algorithm == BLOCK_LOAD_VECTORIZE)
{
if constexpr (detail::is_CacheModifiedInputIterator<RandomAccessIterator>
&& ::cuda::std::is_same_v<detail::it_value_t<RandomAccessIterator>, T>)
if constexpr (detail::is_CacheModifiedInputIterator<RandomAccessIterator>)
{
InternalLoadDirectBlockedVectorized<RandomAccessIterator::__modifier>(linear_tid, block_src_it.ptr, dst_items);
LoadDirectBlockedVectorized<RandomAccessIterator::__modifier>(linear_tid, block_src_it.ptr, dst_items);
}
else if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v<RandomAccessIterator>
&& ::cuda::std::is_same_v<detail::it_value_t<RandomAccessIterator>, T>
// handles contiguous iterators wrapping volatile pointers, like span<volatile int>::iterator
&& !::cuda::std::is_volatile_v<
::cuda::std::remove_reference_t<detail::it_reference_t<RandomAccessIterator>>>)
else if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v<RandomAccessIterator>)
{
InternalLoadDirectBlockedVectorized<LOAD_DEFAULT>(linear_tid, ::cuda::std::to_address(block_src_it), dst_items);
LoadDirectBlockedVectorized(linear_tid, ::cuda::std::to_address(block_src_it), dst_items);
}
else
{
Expand Down
9 changes: 6 additions & 3 deletions cub/cub/device/dispatch/dispatch_topk.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,8 @@ __launch_bounds__(int(current_policy<PolicySelector>().threads_per_block))
policy.items_per_thread,
policy.bits_per_pass,
policy.load_algorithm,
policy.scan_algorithm>;
policy.scan_algorithm,
policy.vec_size>;
using agent_topk_t =
AgentTopK<agent_topk_policy_t,
KeyInputIteratorT,
Expand Down Expand Up @@ -144,7 +145,8 @@ __launch_bounds__(int(current_policy<PolicySelector>().threads_per_block))
policy.items_per_thread,
policy.bits_per_pass,
policy.load_algorithm,
policy.scan_algorithm>;
policy.scan_algorithm,
policy.vec_size>;
using identify_candidates_op_t = NullType;
using agent_topk_t =
AgentTopK<agent_topk_policy_t,
Expand Down Expand Up @@ -205,7 +207,8 @@ __launch_bounds__(int(current_policy<PolicySelector>().threads_per_block))
policy.items_per_thread,
policy.bits_per_pass,
policy.load_algorithm,
policy.scan_algorithm>;
policy.scan_algorithm,
policy.last_filter_vec_size>;
using extract_bin_op_t = NullType;
using agent_topk_t =
AgentTopK<agent_topk_policy_t,
Expand Down
Loading
Loading