Repository navigation
[CUB] Use device-word vectors in StoreDirectBlockedVectorized - #12193
bernhardmgruber wants to merge 3 commits into
Conversation
|
No actionable comments were generated in the recent review.
|
| Layer / File(s) | Summary |
|---|---|
Device-word staging and validation cub/cub/block/block_store.cuh, cub/test/catch2_test_block_store.cu |
The store selects a vector width from the device-word count and stages values in aligned vector storage. It uses direct scalar stores when the output pointer is misaligned. Vectorization no longer requires default-constructible values. The test checks storing a 12-byte type. |
Assessment against linked issues
| Objective | Addressed | Explanation |
|---|---|---|
Use device-word vectors for StoreDirectBlockedVectorized and remove the default-constructibility guard [#12162] |
✅ |
Possibly related issues
- [CUB] Use device-word vectors in StoreDirectBlockedVectorized like BlockLoad #12162: The PR changes the device-word staging and default-constructibility restriction described in this issue.
Possibly related PRs
- NVIDIA/cccl#12163: This PR removes the default-constructibility guard that PR
#12163added while retaining its other eligibility restrictions. - NVIDIA/cccl#12172: Both PRs remove a default-constructibility requirement from CUB vectorized access paths.
- NVIDIA/cccl#12042: Both PRs address vectorized memory staging in CUB; this PR implements device-word staging for stores.
Suggested reviewers: davebayer, srinivasyadav18
Priority:
Change: Bug fix
Merge Risk: 
Moderate · up to db2de
-
The vectorized store change gives types such as 12-byte structs real vector stores. Types with nontrivial assignment can now behave differently depending on output alignment, so the vector path should be limited to trivially copyable types. The public documentation of vectorization restrictions is also out of date. Resolve both before merging, or explicitly accept them.
Comment @coderabbitai help to get the list of available commands.
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.
🟡 Minor · suggestion: Correct the documented vectorization constraints. · block_store.cuh:149-150
cub/cub/block/block_store.cuh:149-150
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winsuggestion: Correct the documented vectorization constraints.
A full tile of a 12-byte custom
Tcan now select device-word vector stores. This description still says custom types always fall back. The algorithm description repeats that claim at Lines 474-475. Update both descriptions. As per coding guidelines, “The@brief/@param/@returndescription must accurately reflect the current functionality of the function.”Source: Coding guidelines
🧹 Nitpick comments: 1
cub/test/catch2_test_block_store.cu: 1
263-263: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winimportant: Include a guaranteed full-tile case.
Ten random sizes do not guarantee
h_input.size() == tile_size. The kernel calls the unguardedStoreonly for a full tile. Partial sizes call the guarded overload, which does not callStoreDirectBlockedVectorized. Add an explicittile_sizecase so theBLOCK_STORE_VECTORIZEinstantiation can cover the changed path.
ℹ️ Review info
⚙️ Run configuration
- Source mode: Auto
- Repository: NVIDIA/cccl
- Branch: main
- File: .coderabbit.yaml
- Commit: 329c4e9
- Layers: Repository YAML
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
7782fd10-ab1d-4001-9a64-6b9329f45387
📒 Files selected for processing: 2
cub/cub/block/block_store.cuhcub/test/catch2_test_block_store.cu
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 6 remain after this review.
|
|
||
| // Direct-store using vector types | ||
| StoreDirectBlocked(linear_tid, block_ptr_vectors, raw_vector); | ||
| StoreDirectBlocked(linear_tid, reinterpret_cast<vector_t*>(block_ptr), vec_items); |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🔎 Supported by static analysis
🏁 Script executed:
git show 329c4e942a0e1cc2605502974f266b9297867420:cub/cub/block/block_store.cuh | sed -n 160,200p
git show 329c4e942a0e1cc2605502974f266b9297867420:cub/cub/block/block_store.cuh | sed -n 835,860p
sed -n 835,855p cub/cub/block/block_store.cuhRepository: NVIDIA/cccl
Length of output: 3753
🏁 Script executed:
set -o pipefail
for rev in 329c4e942a0e1cc2605502974f266b9297867420 babc39abca27022ddad3aa6791e8330421a49ba4; do
echo "===== $rev block_store helper/dispatch ====="
git show "$rev:cub/cub/block/block_store.cuh" | nl -ba | sed -n '155,235p;825,875p'
done
echo "===== load helper references ====="
rg -n -F --glob '*.cuh' -e 'InternalLoadDirectBlockedVectorized' -e 'StoreDirectBlockedVectorized' cub
load_file=$(rg -l -F --glob '*.cuh' 'InternalLoadDirectBlockedVectorized' cub | head -n 1)
if test -n "$load_file"; then
echo "===== $load_file ====="
nl -ba "$load_file" | sed -n '1,260p'
fiRepository: NVIDIA/cccl
Length of output: 24249
🏁 Script executed:
set -o pipefail
echo "===== block load helper ====="
nl -ba cub/cub/block/block_load.cuh | sed -n '175,245p'
echo "===== block load direct/vector callers ====="
nl -ba cub/cub/block/block_load.cuh | sed -n '255,290p;990,1040p'
echo "===== direct store helper ====="
rg -n -F --glob '*.cuh' -e 'void StoreDirectBlocked(' -e 'StoreDirectBlocked(' cub/cub | head -n 20Repository: NVIDIA/cccl
Length of output: 8386
🏁 Script executed:
set -o pipefail
echo "===== StoreDirectBlocked definition ====="
rg -n -F --glob 'block_store.cuh' 'StoreDirectBlocked(' cub/cub/block/block_store.cuh
nl -ba cub/cub/block/block_store.cuh | sed -n '90,155p'
echo "===== InternalLoadDirectBlockedVectorized ====="
nl -ba cub/cub/block/block_load.cuh | sed -n '185,285p'Repository: NVIDIA/cccl
Length of output: 7508
important: Preserve assignment semantics for non-default-constructible T.
The default-constructible case existed before this PR. Removing the guard newly enables the representation-based path for non-default-constructible types. For aligned output, this path stores staged object representations instead of calling T::operator=, while the scalar fallback calls assignment. A copy-constructible, non-default-constructible T can therefore produce different results based on output alignment.
Restrict this path to trivially copyable T.
Suggested fix
&& !::cuda::std::is_volatile_v<
- ::cuda::std::remove_reference_t<detail::it_reference_t<OutputIteratorT>>>)
+ ::cuda::std::remove_reference_t<detail::it_reference_t<OutputIteratorT>>>
+ && ::cuda::std::is_trivially_copyable_v<T>)
This comment has been minimized.
This comment has been minimized.
Stage the items in vectors of UnitWord<T>::DeviceWord, like InternalLoadDirectBlockedVectorized does. This no longer requires T to be default constructible, so drop that condition from BlockStore, and allows real vector stores for types like a 12-byte struct. Fixes: NVIDIA#12162 Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
babc39a to
0c60369
Compare
🥳 CI Workflow Results🟩 Finished in 3h 43m: Pass: 100%/210 | Total: 9d 09h | Max: 3h 20m | Hits: 29%/934985See results here. |
| for (int i = 0; i < ItemsPerThread; i++) | ||
| { | ||
| raw_items[ITEM] = items[ITEM]; | ||
| ::new (reinterpret_cast<T*>(vec_items) + i) T(items[i]); |
There was a problem hiding this comment.
I believe this is not necessary, we should only ever use vectorization for types that support cuda::is_trivially_copyable_v so there is no difference here
There was a problem hiding this comment.
CUB allows store vectorization of types that are not trivially copyable. But we should revert to what is there before.
…ew in StoreDirectBlockedVectorized Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to GitHub limitations.

Minor · important: Align the staging array to alignof(T). · block_store.cuh:182-189
cub/cub/block/block_store.cuh:182-189
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winimportant: Align the staging array to
alignof(T).For a
Twithalignof(T)greater thanalignof(device_word_t),vec_itemshas less alignment than theTobjects accessed throughreinterpret_cast<T*>(vec_items). This breaks the C++ alignment requirement forT. Declare the staging array withalignof(T):Suggested fix
- vector_t vec_items[vectors_per_thread]; + alignas(alignof(T)) vector_t vec_items[vectors_per_thread];This change fixes only the alignment of the staging storage. It does not address the lifetime of the
Tobjects written through the cast.

Review info
-
Run configuration-
- Source mode: Auto
- Repository: NVIDIA/cccl
- Branch: main
- File: .coderabbit.yaml
- Commit: b94ca2a
- Layers: Repository YAML
- Review profile: CHILL
- Plan: Enterprise
- Run ID:
e44d9910-ebda-48b4-9f82-261f1e17eefb
Files selected for processing: 1-
cub/cub/block/block_store.cuh
Files skipped from review as they are similar to previous changes: 1-
- cub/cub/block/block_store.cuh
Included review availability: This review used your included allowance. Your plan provides up to 12 included reviews per hour; 11 remain after this review.
-
…lignof(T) Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
Stages the items in vectors of
UnitWord<T>::DeviceWord, likeInternalLoadDirectBlockedVectorizeddoes. This no longer requiresTto be default constructible, so the condition added in #12163 is dropped fromBlockStore, and types like a 12-byte struct now get real vector stores.I expect some SASS changes here, so let's see what the CI says.
Fixes: #12162
🤖 Generated with Claude Code