Integrate Reduce::ByKey to prefetching#9698
Closed
gonidelis wants to merge 27 commits into
Closed
Conversation
- Add `enum class BlockLoadPrefetch { none, l2, l1 }` to `BlockLoad`
- Add `Prefetch` template param to `BlockLoad` (last position, default `none`, backward-compatible)
- Implement prefetch via `cuda::apply_access_property` for `BLOCK_LOAD_DIRECT`; static_assert on all other algorithms
- Remove `BLOCK_LOAD_DIRECT_PREFETCH` algorithm enum value
- Wire `load_prefetch` through `ReduceByKeyPolicy` -> dispatch -> `agent_reduce_by_key_policy`
- Fire early prefetch in `ConsumeRange` before look-back stall
- All new paths compile to no-op by default - no SASS change
… positional initializers
- use raw ptx intrinsic
Contributor
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
gonidelis
force-pushed
the
prefetch_reducebykey
branch
from
a6c7371 to
8e46ea3
Compare
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughAdds a ChangesBlockLoad prefetch feature
Assessment against linked issues
Suggested reviewers: Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
cub/test/catch2_test_block_load.cu (1)
236-261: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
prefetch_kindomitsBlockLoadPrefetch::l1.The "works across all algorithms and cache levels" test only covers
{l2, bulk_l2}, leaving thel1code path inprefetch_block_load_tile(which emitsprefetch.global.L1) completely untested.✅ Add l1 coverage
using prefetch_kind = c2h::enum_type_list<cub::detail::BlockLoadPrefetch, cub::detail::BlockLoadPrefetch::l2, + cub::detail::BlockLoadPrefetch::l1, cub::detail::BlockLoadPrefetch::bulk_l2>;
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 98daa5dc-f11d-4b23-b857-1b72a6a3d7cb
📒 Files selected for processing (7)
cub/benchmarks/bench/reduce/by_key.cucub/cub/agent/agent_reduce_by_key.cuhcub/cub/block/block_load.cuhcub/cub/device/dispatch/dispatch_reduce_by_key.cuhcub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuhcub/test/catch2_test_block_load.cucub/test/catch2_test_device_reduce_env.cu
Comment on lines
+745
to
+753
| // Fire prefetch before look-back stall so data arrives in L2 in time | ||
| if constexpr (AgentReduceByKeyPolicyT::LOAD_PREFETCH != detail::BlockLoadPrefetch::none) | ||
| { | ||
| const int items_to_prefetch = static_cast<int>(::cuda::std::min(num_remaining, static_cast<OffsetT>(TILE_ITEMS))); | ||
| detail::prefetch_block_load_tile<BLOCK_THREADS, 128, AgentReduceByKeyPolicyT::LOAD_PREFETCH>( | ||
| threadIdx.x, d_keys_in + tile_offset, items_to_prefetch); | ||
| detail::prefetch_block_load_tile<BLOCK_THREADS, 128, AgentReduceByKeyPolicyT::LOAD_PREFETCH>( | ||
| threadIdx.x, d_values_in + tile_offset, items_to_prefetch); | ||
| } |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm CacheModifiedInputIterator's iterator_concept / to_address support
rg -n "class CacheModifiedInputIterator" -A 80 cub/cub/iterator/cache_modified_input_iterator.cuh | rg -n "iterator_concept|iterator_category|to_address|contiguous"Repository: NVIDIA/cccl
Length of output: 260
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant files and read the smallest useful slices.
for f in \
cub/cub/iterator/cache_modified_input_iterator.cuh \
cub/cub/block/block_load.cuh \
cub/cub/agent/agent_reduce_by_key.cuh \
cub/cub/agent/agent_reduce_by_key.tuning.cuh \
cub/cub/iterator/util_type.cuh
do
echo "=== $f ==="
wc -l "$f"
done
echo
echo "=== cache_modified_input_iterator.cuh excerpt ==="
sed -n '1,220p' cub/cub/iterator/cache_modified_input_iterator.cuh
echo
echo "=== block_load.cuh prefetch-related excerpt ==="
rg -n "prefetch|can_prefetch|contiguous_iterator|CacheModifiedInputIterator" -A 8 -B 8 cub/cub/block/block_load.cuh
echo
echo "=== agent_reduce_by_key.cuh wrapping/prefetch excerpt ==="
sed -n '180,220p' cub/cub/agent/agent_reduce_by_key.cuh
echo
sed -n '730,760p' cub/cub/agent/agent_reduce_by_key.cuh
echo
echo "=== tuning references ==="
rg -n "load_prefetch|BlockLoadPrefetch::l2|LOAD_MODIFIER|LOAD_DEFAULT" -A 4 -B 4 cub/cub/agent cub/cub | head -n 200Repository: NVIDIA/cccl
Length of output: 551
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "=== locate tuning file ==="
fd -a 'agent_reduce_by_key' cub/cub/agent
echo
echo "=== cache_modified_input_iterator.cuh relevant lines ==="
rg -n "class CacheModifiedInputIterator|iterator_category|iterator_concept|to_address|contiguous_iterator|operator\\+|operator\\*" -A 6 -B 6 cub/cub/iterator/cache_modified_input_iterator.cuh
echo
echo "=== block_load.cuh prefetch and iterator constraints ==="
rg -n "prefetch|can_prefetch_block_load|contiguous_iterator|CacheModifiedInputIterator|load_modifier" -A 8 -B 8 cub/cub/block/block_load.cuh
echo
echo "=== agent_reduce_by_key.cuh wrapping and prefetch sites ==="
sed -n '185,215p' cub/cub/agent/agent_reduce_by_key.cuh
echo
sed -n '740,758p' cub/cub/agent/agent_reduce_by_key.cuhRepository: NVIDIA/cccl
Length of output: 21012
important: prefetch needs the unwrapped input iterator prefetch_block_load_tile is suppressed for CacheModifiedInputIterator, so calling it on d_keys_in + tile_offset / d_values_in + tile_offset is a no-op for raw-pointer inputs. That leaves LOAD_PREFETCH ineffective on the common pointer-based DeviceReduce::ReduceByKey path; prefetch the raw pointer before wrapping, or use a separate unwrapped prefetch path.
1 task
Contributor
😬 CI Workflow Results🟥 Finished in 2h 42m: Pass: 96%/287 | Total: 12d 22h | Max: 2h 41m | Hits: 19%/1087839See results here. |
Member
Author
|
I am closing this in favor of the new design #9723 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Builds on top of #9496 and integrated cub reduce_by_key to prefetching
Full scan b200 tuning currently running
fixes #9566