Skip to content

[CUB] Take a block offset in LoadDirectStripedVectorized #12601

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
base: main
Choose a base branch
from
+49 −26
Open

[CUB] Take a block offset in LoadDirectStripedVectorized #12601

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
47 changes: 25 additions & 22 deletions cub/cub/block/block_load.cuh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -396,7 +396,8 @@ load_direct_striped_vectorized_aligned(int linear_tid, const InputT* block_src_p
} // namespace detail

//! @rst
//! Load a linear segment of items into a striped arrangement of vectors across the thread block.
//! Load a linear segment of items into a striped arrangement of vectors across the thread block, from ``block_offset``
//! items after ``block_src_ptr``.
//!
//! .. versionadded:: 3.6.0
//!
Expand All @@ -405,8 +406,9 @@ load_direct_striped_vectorized_aligned(int linear_tid, const InputT* block_src_p
//! (*i* + (*block-threads* * (*items-per-thread* / ``VecSize`` - 1))). With ``VecSize == 1``, this is the arrangement
//! of ``LoadDirectStriped``.
//!
//! Each vector is loaded as a ``CubVector<InputT, VecSize>::Type``. The following conditions will prevent
//! vectorization and the same arrangement is loaded item by item:
//! Each vector is loaded as a ``CubVector<InputT, VecSize>::Type``. ``block_offset`` must be a multiple of ``VecSize``,
//! so that the vectors have the alignment of ``block_src_ptr``. The following conditions will prevent vectorization and
//! the same arrangement is loaded item by item:
//!
//! - ``block_src_ptr`` is not aligned to the vector access width
//! - ``InputT`` is volatile or not trivially copyable
Expand All @@ -431,12 +433,18 @@ load_direct_striped_vectorized_aligned(int linear_tid, const InputT* block_src_p
//! @tparam ItemsPerThread
//! **[inferred]** The number of items partitioned onto each thread.
//!
//! @tparam OffsetT
//! **[inferred]** The integral type of ``block_offset``
//!
//! @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
//! The thread block's base pointer for loading from
//! The base pointer whose alignment is tested
//!
//! @param[in] block_offset
//! The offset of the thread block's vectors from ``block_src_ptr``
//!
//! @param[out] dst_items
//! Destination to load data into
Expand All @@ -445,13 +453,17 @@ template <int ThreadsPerBlock,
CacheLoadModifier Modifier = LOAD_DEFAULT,
typename InputT,
typename T,
int ItemsPerThread>
_CCCL_DEVICE _CCCL_FORCEINLINE void
LoadDirectStripedVectorized(int linear_tid, const InputT* block_src_ptr, T (&dst_items)[ItemsPerThread])
int ItemsPerThread,
typename OffsetT>
_CCCL_DEVICE _CCCL_FORCEINLINE void LoadDirectStripedVectorized(
int linear_tid, const InputT* block_src_ptr, OffsetT block_offset, T (&dst_items)[ItemsPerThread])
Comment on lines +458 to +459

@bernhardmgruber bernhardmgruber Nov 11, 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.

As discussed offline, I think this shouldn't be necessary since something like LoadDirectStripedVectorized(tid, assume_aligned<A>(ptr + offset), items) should give the desired effect.

I am especially unhappy because this is a public API and would then be the only one with an offset argument. Users will be confused.

If we really need this, can we make the one with an offset parameter internal and leave the public one as is?

{
static_assert(ItemsPerThread % VecSize == 0, "VecSize must evenly divide ItemsPerThread");
constexpr int vec_size = static_cast<int>(VecSize);

// The vectors start a whole number of vectors after block_src_ptr, thus they have its alignment
const InputT* const tile_src_ptr = block_src_ptr + block_offset;

// TODO(bgruber): remove the check for is_primitive<InputT> in CCCL 4.0
constexpr bool can_vectorize =
!::cuda::std::is_volatile_v<InputT> && (::cuda::is_trivially_copyable_v<InputT> || detail::is_primitive_v<InputT>);
Expand All @@ -460,27 +472,18 @@ LoadDirectStripedVectorized(int linear_tid, const InputT* block_src_ptr, T (&dst
if (::cuda::std::is_sufficiently_aligned<alignof(typename CubVector<InputT, vec_size>::Type)>(block_src_ptr))
{
_CUB::detail::load_direct_striped_vectorized_aligned<ThreadsPerBlock, VecSize, Modifier>(
linear_tid, block_src_ptr, dst_items);
linear_tid, tile_src_ptr, dst_items);
return;
}
}

// Load the same arrangement item by item
const InputT* thread_src_ptr = block_src_ptr + static_cast<unsigned int>(linear_tid) * vec_size;
const InputT* thread_src_ptr = tile_src_ptr + static_cast<unsigned int>(linear_tid) * vec_size;
_CCCL_PRAGMA_UNROLL_FULL()
for (int i = 0; i < ItemsPerThread / vec_size; i++)
{
if constexpr (Modifier != LOAD_DEFAULT)
{
_CUB::LoadDirectBlocked(i * ThreadsPerBlock,
CacheModifiedInputIterator<Modifier, InputT, int>(thread_src_ptr),
reinterpret_cast<T(&)[vec_size]>(dst_items[i * vec_size]));
}
else
{
_CUB::LoadDirectBlocked(
i * ThreadsPerBlock, thread_src_ptr, reinterpret_cast<T(&)[vec_size]>(dst_items[i * vec_size]));
}
_CUB::LoadDirectBlocked(
i * ThreadsPerBlock, thread_src_ptr, reinterpret_cast<T(&)[vec_size]>(dst_items[i * vec_size]));
Comment on lines +485 to +486

@coderabbitai coderabbitai Bot Nov 11, 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.

Major · Functional Correctness · Quick win

Supported by static analysis

Script executed:

git diff cb60424c559952f127b9bdb12afeb4b40e2b3271 f41c362208b10611684300eb522e5eadad37318f -- cub/cub/block/block_load.cuh
sed -n '440,495p' cub/cub/block/block_load.cuh
sed -n '1225,1260p' cub/cub/block/block_load.cuh

Repository: NVIDIA/cccl

Length of output: 9864

Script executed:

rg -n -F -e 'LoadDirectBlocked(' -e 'LoadDirectStripedVectorized(' -- cub/cub/block/block_load.cuh
rg -n -F -e 'LOAD_VOLATILE' -e 'ThreadLoad(' -e 'CacheModifiedInputIterator' -- cub/cub/thread/thread_load.cuh cub/cub/iterator/cache_modified_input_iterator.cuh cub/cub/util_ptx.cuh cub/cub/util_type.cuh
cat -n cub/cub/block/block_load.cuh | sed -n '55,145p;345,405p;405,490p;1190,1265p'
cat -n cub/cub/iterator/cache_modified_input_iterator.cuh | sed -n '55,250p'
cat -n cub/cub/thread/thread_load.cuh | sed -n '45,240p'

Repository: NVIDIA/cccl

Length of output: 38247

Script executed:

cat -n cub/cub/thread/thread_load.cuh | sed -n '240,400p'

Repository: NVIDIA/cccl

Length of output: 8139

important: Preserve Modifier in scalar loads.

When a non-volatile input falls back to scalar loads, LoadDirectBlocked uses ordinary pointer loads and ignores Modifier. This also affects BlockLoad with a cache-modified iterator. For LOAD_VOLATILE, the loads lack the volatile semantics of ThreadLoad<LOAD_VOLATILE>.

Preserve the modifier in the fallback and test the generated loads. Preserve the supported input types and their copy behavior.

}
Comment on lines 483 to 487

@miscco miscco Nov 11, 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.

This is a bug compatible change to keep this as BlockLoadDirectVectorized, which drops the load modifier on fallback

}

Expand Down Expand Up @@ -1237,13 +1240,13 @@ public:
&& ::cuda::std::is_same_v<detail::it_value_t<RandomAccessIterator>, T>)
{
_CUB::LoadDirectStripedVectorized<ThreadsPerBlock, vec_size, RandomAccessIterator::__modifier>(
linear_tid, block_src_it.ptr, dst_items);
linear_tid, block_src_it.ptr, 0, dst_items);
}
else if constexpr (THRUST_NS_QUALIFIER::is_contiguous_iterator_v<RandomAccessIterator>
&& ::cuda::std::is_same_v<detail::it_value_t<RandomAccessIterator>, T>)
{
_CUB::LoadDirectStripedVectorized<ThreadsPerBlock, vec_size, LOAD_DEFAULT>(
linear_tid, ::cuda::std::to_address(block_src_it), dst_items);
linear_tid, ::cuda::std::to_address(block_src_it), 0, dst_items);
}
else
{
Expand Down
28 changes: 24 additions & 4 deletions cub/test/catch2_test_block_load.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -320,26 +320,28 @@ CUB_TEST("Vectorized block load converts from iterators of a different value typ
}

template <int ThreadsInBlock, int VecSize, cub::CacheLoadModifier Modifier, int ItemsPerThread, typename InputT, typename T>
__global__ void load_direct_striped_vectorized_kernel(const InputT* input, T* output)
__global__ void load_direct_striped_vectorized_kernel(const InputT* input, int block_offset, T* output)
{
T data[ItemsPerThread];
cub::LoadDirectStripedVectorized<ThreadsInBlock, VecSize, Modifier>(static_cast<int>(threadIdx.x), input, data);
cub::LoadDirectStripedVectorized<ThreadsInBlock, VecSize, Modifier>(
static_cast<int>(threadIdx.x), input, block_offset, data);
for (int i = 0; i < ItemsPerThread; i++)
{
output[(threadIdx.x + ThreadsInBlock * (i / VecSize)) * VecSize + i % VecSize] = data[i];
}
}

template <int VecSize, cub::CacheLoadModifier Modifier = cub::CacheLoadModifier::LOAD_DEFAULT, typename T, typename InputT>
void test_load_direct_striped_vectorized(const c2h::device_vector<T>& d_expected, const InputT* input)
void test_load_direct_striped_vectorized(
const c2h::device_vector<T>& d_expected, const InputT* input, int block_offset = 0)
{
constexpr int items_per_thread = 8;
constexpr int threads_in_block = 64;
REQUIRE(d_expected.size() == items_per_thread * threads_in_block);

c2h::device_vector<T> d_output(d_expected.size(), thrust::no_init);
load_direct_striped_vectorized_kernel<threads_in_block, VecSize, Modifier, items_per_thread>
<<<1, threads_in_block>>>(input, thrust::raw_pointer_cast(d_output.data()));
<<<1, threads_in_block>>>(input, block_offset, thrust::raw_pointer_cast(d_output.data()));
REQUIRE(cudaSuccess == cudaPeekAtLastError());
REQUIRE(cudaSuccess == cudaDeviceSynchronize());
REQUIRE(d_expected == d_output);
Expand Down Expand Up @@ -367,6 +369,24 @@ CUB_TEST("LoadDirectStripedVectorized loads a striped arrangement of vectors of
test_load_direct_striped_vectorized<4, cub::CacheLoadModifier::LOAD_LDG>(d_expected, input);
}

// block_offset is a multiple of the vectors, thus the vectors keep the alignment of the input
CUB_TEST(
"LoadDirectStripedVectorized loads the vectors block_offset items after the input", "[load][block]", CUB_SMALL, types)
{
using type = c2h::get<0, TestType>;
constexpr int tile_size = 8 * 64;

const int offset = GENERATE(0, 1);
const int block_offset = GENERATE_COPY(8, 3 * tile_size);
c2h::device_vector<type> d_input(offset + block_offset + tile_size);
c2h::gen(C2H_SEED(1), d_input);
const c2h::device_vector<type> d_expected(d_input.begin() + offset + block_offset, d_input.end());
const type* input = thrust::raw_pointer_cast(d_input.data()) + offset;

test_load_direct_striped_vectorized<4>(d_expected, input, block_offset);
test_load_direct_striped_vectorized<4, cub::CacheLoadModifier::LOAD_LDG>(d_expected, input, block_offset);
}

CUB_TEST_CASE("LoadDirectStripedVectorized converts the items to the destination type", "[load][block]", CUB_SMALL)
{
constexpr int tile_size = 8 * 64;
Expand Down
Loading