Implement prefetching#9723
Open
gonidelis wants to merge 4 commits into
Open
Conversation
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. |
4 tasks
Member
Author
|
See: #3126 (comment) that justifies this PR |
* Use 32-bit multiplication for the tile byte count (truncated anyway) * Guard the entire bulk_l2 branch with NV_IF_TARGET(NV_PROVIDES_SM_90) instead of only the asm statement * Handle LoadPrefetch::none inside prefetch_tile so call sites need no enablement check * Move the Prefetch template parameter before the defaulted PrefetchStride so common call sites don't restate the stride
Contributor
📝 WalkthroughSummary by CodeRabbit
WalkthroughThis PR adds a ChangesReduce-by-key prefetch support
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
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 54e03265-c405-4ed7-9a6a-2dda8d2ab5e1
📒 Files selected for processing (6)
cub/benchmarks/bench/reduce/by_key.cucub/cub/agent/agent_reduce_by_key.cuhcub/cub/detail/prefetch.cuhcub/cub/device/dispatch/dispatch_reduce_by_key.cuhcub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuhcub/test/catch2_test_device_reduce_by_key_env.cu
Comment on lines
+746
to
+772
| // Fire prefetch hints before the decoupled look-back stall inside ConsumeTile, so the | ||
| // stall hides the L2 fetch latency and the tile's data has arrived by the time it is loaded. | ||
| // prefetch_tile is a no-op for LoadPrefetch::none, so no enablement check is needed here. | ||
| // The wrapped iterators apply the policy's LOAD_MODIFIER to the user's raw pointers; unlike a | ||
| // user-supplied CacheModifiedInputIterator, this internal wrapper expresses no intent to bypass | ||
| // caches, so unwrapping it for the prefetch hint is legitimate. | ||
| const int items_to_prefetch = static_cast<int>(::cuda::std::min(num_remaining, static_cast<OffsetT>(TILE_ITEMS))); | ||
| if constexpr (::cuda::std::is_pointer_v<KeysInputIteratorT>) | ||
| { | ||
| detail::prefetch_tile<BLOCK_THREADS, AgentReduceByKeyPolicyT::LOAD_PREFETCH>( | ||
| threadIdx.x, (d_keys_in + tile_offset).ptr, items_to_prefetch); | ||
| } | ||
| else | ||
| { | ||
| detail::prefetch_tile<BLOCK_THREADS, AgentReduceByKeyPolicyT::LOAD_PREFETCH>( | ||
| threadIdx.x, d_keys_in + tile_offset, items_to_prefetch); | ||
| } | ||
| if constexpr (::cuda::std::is_pointer_v<ValuesInputIteratorT>) | ||
| { | ||
| detail::prefetch_tile<BLOCK_THREADS, AgentReduceByKeyPolicyT::LOAD_PREFETCH>( | ||
| threadIdx.x, (d_values_in + tile_offset).ptr, items_to_prefetch); | ||
| } | ||
| else | ||
| { | ||
| detail::prefetch_tile<BLOCK_THREADS, AgentReduceByKeyPolicyT::LOAD_PREFETCH>( | ||
| threadIdx.x, d_values_in + tile_offset, items_to_prefetch); | ||
| } |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the relevant area in cub/cub/agent/agent_reduce_by_key.cuh
file="cub/cub/agent/agent_reduce_by_key.cuh"
wc -l "$file"
sed -n '700,820p' "$file"
# Find prefetch_tile definitions/usages
rg -n "prefetch_tile<|prefetch_tile\s*\(" cub/cub -S
# Inspect likely definition file(s) if found
fd -a "prefetch" cub/cubRepository: NVIDIA/cccl
Length of output: 4238
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Read the implementation of detail::prefetch_tile if located by search.
rg -n "template <.*prefetch_tile|prefetch_tile\(" cub -S --glob '!**/build/**'Repository: NVIDIA/cccl
Length of output: 395
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Narrow to the specific header/namespace that defines detail::prefetch_tile.
rg -n "namespace detail|prefetch_tile" cub/cub -S --glob '*.cuh' --glob '*.cu' --glob '*.h' --glob '*.hpp'Repository: NVIDIA/cccl
Length of output: 35907
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Show the exact surrounding control flow in agent_reduce_by_key.cuh.
file="cub/cub/agent/agent_reduce_by_key.cuh"
sed -n '740,810p' "$file"
# Locate detail::prefetch_tile implementation and its handling of item counts.
rg -n "prefetch_tile" cub/cub -S --glob '*.cuh' --glob '*.cu' --glob '*.h' --glob '*.hpp' -A 25 -B 10Repository: NVIDIA/cccl
Length of output: 17098
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Find the actual definition of detail::prefetch_tile.
rg -n "detail::prefetch_tile|prefetch_tile" . -S --glob '*.cuh' --glob '*.cu' --glob '*.h' --glob '*.hpp' --glob '*.cc' --glob '*.cpp' -A 30 -B 10Repository: NVIDIA/cccl
Length of output: 16467
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect how tiles/blocks are launched and whether num_remaining can be non-positive.
rg -n "max_dim_x|num_tiles|num_remaining|ConsumeRange|tile_offset" cub/cub/agent/agent_reduce_by_key.cuh -A 20 -B 20Repository: NVIDIA/cccl
Length of output: 11660
🏁 Script executed:
#!/bin/bash
set -euo pipefail
sed -n '88,134p' cub/cub/detail/prefetch.cuhRepository: NVIDIA/cccl
Length of output: 2512
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Check other call sites to see whether negative item counts are expected anywhere else.
rg -n "prefetch_tile<.*\(|prefetch_tile\(" cub/cub -S --glob '*.cuh' --glob '*.cu' --glob '*.h' --glob '*.hpp' -A 8 -B 8Repository: NVIDIA/cccl
Length of output: 7888
important: gate the prefetch hints on num_remaining > 0. prefetch_tile casts the count to unsigned, so a block that lands past the end turns a negative num_remaining into a huge byte request and can issue a runaway prefetch before the existing ConsumeTile guard skips the tile.
Contributor
😬 CI Workflow Results🟥 Finished in 2h 35m: Pass: 94%/287 | Total: 11d 15h | Max: 2h 34m | Hits: 18%/1027667See results here. |
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.
fixes #9616 and #9698 by taking a completely different approach.
Introduces prefetching as a standalone facility: a
detail::LoadPrefetch enum+detail::prefetch_tile()free function. Also shows MVP wiring of prefetching intoDeviceReduce::ReduceByKeyas a tuning-policy knob (ReduceByKeyPolicy::load_prefetch).Prefecthing should not be part of
BlockLoadas a template parameter. Prefetch hints only pay off with lead time. Emitting insideLoad()fires the hint instructions immediately before the loads of the same addresses. Firing the same hints earlier, in the agent's schedule (before reduce_by_key's look-back machinery), wins up to −14% per CT workload.