From f41c362208b10611684300eb522e5eadad37318f Mon Sep 17 00:00:00 2001 From: Michael Schellenberger Costa Date: Wed, 11 Nov 2026 15:54:39 +0100 Subject: [PATCH] [CUB] Take a block offset in `LoadDirectStripedVectorized` The function is new in 3.6, so its first signature loads the vectors `block_offset` items after the base pointer. The offset is a multiple of the vector size, which keeps the vectors aligned with that pointer, and later changes can select the vector width without adding an overload. --- cub/cub/block/block_load.cuh | 47 ++++++++++++++++-------------- cub/test/catch2_test_block_load.cu | 28 +++++++++++++++--- 2 files changed, 49 insertions(+), 26 deletions(-) diff --git a/cub/cub/block/block_load.cuh b/cub/cub/block/block_load.cuh index 9a05cd8d80ab..79853a149cb2 100644 --- a/cub/cub/block/block_load.cuh +++ b/cub/cub/block/block_load.cuh @@ -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 //! @@ -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::Type``. The following conditions will prevent -//! vectorization and the same arrangement is loaded item by item: +//! Each vector is loaded as a ``CubVector::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 @@ -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 @@ -445,13 +453,17 @@ template -_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]) { static_assert(ItemsPerThread % VecSize == 0, "VecSize must evenly divide ItemsPerThread"); constexpr int vec_size = static_cast(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 in CCCL 4.0 constexpr bool can_vectorize = !::cuda::std::is_volatile_v && (::cuda::is_trivially_copyable_v || detail::is_primitive_v); @@ -460,27 +472,18 @@ LoadDirectStripedVectorized(int linear_tid, const InputT* block_src_ptr, T (&dst if (::cuda::std::is_sufficiently_aligned::Type)>(block_src_ptr)) { _CUB::detail::load_direct_striped_vectorized_aligned( - 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(linear_tid) * vec_size; + const InputT* thread_src_ptr = tile_src_ptr + static_cast(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(thread_src_ptr), - reinterpret_cast(dst_items[i * vec_size])); - } - else - { - _CUB::LoadDirectBlocked( - i * ThreadsPerBlock, thread_src_ptr, reinterpret_cast(dst_items[i * vec_size])); - } + _CUB::LoadDirectBlocked( + i * ThreadsPerBlock, thread_src_ptr, reinterpret_cast(dst_items[i * vec_size])); } } @@ -1237,13 +1240,13 @@ public: && ::cuda::std::is_same_v, T>) { _CUB::LoadDirectStripedVectorized( - 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 && ::cuda::std::is_same_v, T>) { _CUB::LoadDirectStripedVectorized( - linear_tid, ::cuda::std::to_address(block_src_it), dst_items); + linear_tid, ::cuda::std::to_address(block_src_it), 0, dst_items); } else { diff --git a/cub/test/catch2_test_block_load.cu b/cub/test/catch2_test_block_load.cu index 8dd630e65763..bad4c71eae3e 100644 --- a/cub/test/catch2_test_block_load.cu +++ b/cub/test/catch2_test_block_load.cu @@ -320,10 +320,11 @@ CUB_TEST("Vectorized block load converts from iterators of a different value typ } template -__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(static_cast(threadIdx.x), input, data); + cub::LoadDirectStripedVectorized( + static_cast(threadIdx.x), input, block_offset, data); for (int i = 0; i < ItemsPerThread; i++) { output[(threadIdx.x + ThreadsInBlock * (i / VecSize)) * VecSize + i % VecSize] = data[i]; @@ -331,7 +332,8 @@ __global__ void load_direct_striped_vectorized_kernel(const InputT* input, T* ou } template -void test_load_direct_striped_vectorized(const c2h::device_vector& d_expected, const InputT* input) +void test_load_direct_striped_vectorized( + const c2h::device_vector& d_expected, const InputT* input, int block_offset = 0) { constexpr int items_per_thread = 8; constexpr int threads_in_block = 64; @@ -339,7 +341,7 @@ void test_load_direct_striped_vectorized(const c2h::device_vector& d_expected c2h::device_vector d_output(d_expected.size(), thrust::no_init); load_direct_striped_vectorized_kernel - <<<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); @@ -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 d_input(offset + block_offset + tile_size); + c2h::gen(C2H_SEED(1), d_input); + const c2h::device_vector 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;