Skip to content

[CUB] Use device-word vectors in StoreDirectBlockedVectorized - #12193

Open
bernhardmgruber wants to merge 3 commits into
NVIDIA:mainfrom
bernhardmgruber:block-store-device-word-vectors
Open

bernhardmgruber wants to merge 3 commits into
NVIDIA:mainfrom
bernhardmgruber:block-store-device-word-vectors

Conversation

@bernhardmgruber

@bernhardmgruber bernhardmgruber commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

Stages the items in vectors of UnitWord<T>::DeviceWord, like InternalLoadDirectBlockedVectorized does. This no longer requires T to be default constructible, so the condition added in #12163 is dropped from BlockStore, 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

@bernhardmgruber
bernhardmgruber requested a review from a team as a code owner October 22, 2026 21:08
@github-project-automation github-project-automation Bot moved this to Todo in CCCL Oct 22, 2026
@cccl-authenticator-app cccl-authenticator-app Bot moved this from Todo to In Review in CCCL Oct 22, 2026
@coderabbitai

coderabbitai Bot commented Oct 22, 2026 •

Copy link
Copy Markdown
Contributor

Review in Change Stack →

No actionable comments were generated in the recent review.

Recent 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: 1d16216a-315b-41c9-b96a-99d5a05094ac
Commits

Reviewing files that changed from the base of the PR and between 0e0a919 and db2de1e.

Files selected for processing: 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.

📝 Summary

Summary by CodeRabbit

  • Bug Fixes
    • Block stores now support eligible value types that cannot be default-constructed.
    • Stores use vectorized writes when the output address meets the required alignment, while retaining scalar writes for misaligned addresses.
    • Improved handling of 12-byte values across supported items-per-thread and store-algorithm configurations, with validation that stored values match the inputs.

BlockStore now stages vectorized stores in device-word vectors. Misaligned output pointers use the existing direct scalar stores. Vectorization no longer requires default-constructible values. A test covers storing a 12-byte value type.

Changes: Device-word vectorized block stores

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

Possibly related PRs

  • NVIDIA/cccl#12163: This PR removes the default-constructibility guard that PR #12163 added 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: Normal

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.

- [ ] Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

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.

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to GitHub limitations.

⚠️ Outside diff range comments (1)

🟡 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 win

suggestion: Correct the documented vectorization constraints.

A full tile of a 12-byte custom T can 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/@return description 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 win

important: Include a guaranteed full-tile case.

Ten random sizes do not guarantee h_input.size() == tile_size. The kernel calls the unguarded Store only for a full tile. Partial sizes call the guarded overload, which does not call StoreDirectBlockedVectorized. Add an explicit tile_size case so the BLOCK_STORE_VECTORIZE instantiation 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
📥 Commits

Reviewing files that changed from the base of the PR and between 329c4e9 and babc39a.

📒 Files selected for processing: 2
  • cub/cub/block/block_store.cuh
  • cub/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.

Comment thread cub/cub/block/block_store.cuh Outdated

// Direct-store using vector types
StoreDirectBlocked(linear_tid, block_ptr_vectors, raw_vector);
StoreDirectBlocked(linear_tid, reinterpret_cast<vector_t*>(block_ptr), vec_items);

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.

🎯 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.cuh

Repository: 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'
fi

Repository: 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 20

Repository: 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>)

@github-actions

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>
@bernhardmgruber
bernhardmgruber force-pushed the block-store-device-word-vectors branch from babc39a to 0c60369 Compare October 23, 2026 13:19

@miscco miscco left a comment

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.

Lets wait on #12042 which brings in unified machinery for vectorized loads which should also apply to stores

@github-project-automation github-project-automation Bot moved this from In Review to In Progress in CCCL Oct 23, 2026
@github-actions

Copy link
Copy Markdown
Contributor

🥳 CI Workflow Results

🟩 Finished in 3h 43m: Pass: 100%/210 | Total: 9d 09h | Max: 3h 20m | Hits: 29%/934985

See results here.

Comment thread cub/cub/block/block_store.cuh Outdated
for (int i = 0; i < ItemsPerThread; i++)
{
raw_items[ITEM] = items[ITEM];
::new (reinterpret_cast<T*>(vec_items) + i) T(items[i]);

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.

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

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.

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>
@bernhardmgruber

Copy link
Copy Markdown
Contributor Author

Lets wait on #12042 which brings in unified machinery for vectorized loads which should also apply to stores

Let's try to get this PR in first, since it is SASS neutral. AFAIU, #12042 will use a different vectorization approach we may have SASS impact.

@github-project-automation github-project-automation Bot moved this from In Progress to In Review in CCCL Oct 28, 2026
@bernhardmgruber
bernhardmgruber enabled auto-merge (squash) October 28, 2026 16:10

@coderabbitai coderabbitai Bot left a comment

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.

Caution

Some comments are outside the diff and can’t be posted inline due to GitHub limitations.

⚠️ Outside diff range comments (1)

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 win

important: Align the staging array to alignof(T).

For a T with alignof(T) greater than alignof(device_word_t), vec_items has less alignment than the T objects accessed through reinterpret_cast<T*>(vec_items). This breaks the C++ alignment requirement for T. Declare the staging array with alignof(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 T objects 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
Commits

Reviewing files that changed from the base of the PR and between 0c60369 and 0e0a919.

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>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: In Review

Development

Successfully merging this pull request may close these issues.

[CUB] Use device-word vectors in StoreDirectBlockedVectorized like BlockLoad

2 participants