Fix DecoupledLookback cross-block coherence (completes #91)#98
Open
shreyas-omkar wants to merge 4 commits into
Open
Fix DecoupledLookback cross-block coherence (completes #91)#98shreyas-omkar wants to merge 4 commits into
shreyas-omkar wants to merge 4 commits into
Conversation
…kback Re-enable DecoupledLookback() coverage in the non-uniform exclusive accumulate test (excluded on the main branch). This branch carries the work-in-progress fix for DecoupledLookback's exclusive carries, so it must exercise that path. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…need The decoupled-lookback aggregate publish/consume protocol is not memory- coherent across blocks (stale L1 aggregate reads + missing device-scope ordering), causing ~40% failures on smaller GPUs. The proper fix needs a native device threadfence; `UnsafeAtomics.fence` and acquire/release atomics do not lower on recent NVPTX toolchains (LLVM 18, sm_80) — only `monotonic` does. Document this inline as the path forward. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The exclusive-carry fix published each block's inclusive aggregate in a
dedicated array, but the publish/consume protocol was not memory-coherent
across blocks: a consumer could observe a fresh flag == ACC_FLAG_A yet read a
stale L1 copy of the aggregate, and the producer's aggregate/flag stores were
not ordered device-scope. This caused ~40% failures on smaller GPUs.
Fix both problems:
- Coherence: read/write the aggregate with a relaxed (monotonic), L1-bypassing
atomic, reinterpreted through a same-width unsigned integer so float element
types are supported (compile-time fallback to a plain access otherwise).
- Ordering: introduce a device-scope `_decoupled_fence()` separating the
aggregate store from the flag store (producer) and the flag load from the
aggregate load (consumer). The generic definition is an acquire-release
atomic fence (lowers on OpenCL/SPIR-V, Metal, oneAPI, AMDGPU); CUDA overrides
it with `CUDA.threadfence()` (membar.gl) via a new package extension, since
the NVPTX backend does not select scoped atomic fences.
`aggregates` is now always its own contiguous array (never a temp view) so
`pointer(aggregates, i)` is valid inside the kernel.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
9ea5555 to
63a1fcc
Compare
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
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.
Completes #91 — a working version of the
DecoupledLookback()scan.Background
#91 corrected
DecoupledLookback()'s exclusive multi-block carries by publishing each block's inclusive aggregate in a dedicated array. That fixed the algorithm, but the cross-block publish/consume protocol was still not memory-coherent, so it failed ~40% of the time on smaller GPUs (e.g. A100 MIG 1g.5gb, sm_80) for both inclusive and exclusive scans.Two problems remained:
flag == ACC_FLAG_Ayet read a stale aggregate.Fix
monotonic), L1-bypassing atomic, reinterpreted through a same-width unsigned integer so float element types are supported (compile-time fallback to a plain access for other widths).There is no vendor-agnostic device-scope fence in the ecosystem today (KA only has work-group
barrier/@synchronize;UnsafeAtomics.fence(acq_rel)does not lower on NVPTX;CUDA.threadfenceis CUDA-only). So_decoupled_fence()composes the already-available primitives, kept self-contained here:UnsafeAtomics.fence(acq_rel)(lowers on OpenCL/SPIR-V, Metal, oneAPI, AMDGPU),CUDA.threadfence()(membar.gl) via a small CUDA package extension, since NVPTX does not select scoped atomic fences.No new hard dependency:
CUDAis a weakdep behind an extension.Also:
aggregatesis now always its own contiguous array (never a temp view) sopointer(aggregates, i)is valid inside the kernel; the non-uniform exclusive-scan test is re-extended to coverDecoupledLookback().Validation
Performance / when to use it
DecoupledLookback()is now correct, but it is not the fastest scan on CUDA — launches are cheap there, so the standaloneScanPrefixes()(which stays the default on every backend) wins, and the gap grows withn:So this PR does not change the default (
ScanPrefixes()everywhere). Its value is makingDecoupledLookback()correct so it can be opted into on dispatch-bound backends (e.g. Metal), where collapsing kernel launches is the win — the motivation for the single-pass/onesweep direction.