From 63aacc181c46c7b6747d70efd342ff4afabe1aec Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 30 Jun 2026 11:59:44 -0400 Subject: [PATCH 1/3] test(scan): re-extend non-uniform exclusive scan test to DecoupledLookback 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) --- test/generic/accumulate.jl | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/test/generic/accumulate.jl b/test/generic/accumulate.jl index 0759308..d308f5a 100644 --- a/test/generic/accumulate.jl +++ b/test/generic/accumulate.jl @@ -34,19 +34,14 @@ ALGS = AK.AccumulateAlgorithm[AK.ScanPrefixes()] end # Non-uniform data exposes ScanPrefixes block-carry bugs that all-ones data masks. - # NOTE: DecoupledLookback() is excluded here: its exclusive multi-block carries are still - # incorrect on non-uniform data (the all-ones exclusive tests above mask that bug). Fixing it - # requires a coherent cross-block aggregate publish that is tracked separately. - if alg isa AK.ScanPrefixes - for _ in 1:200 - num_elems = rand(513:100_000) - block_size = rand([16, 32, 64, 128, 256]) - init = rand(Int32(-100):Int32(100)) - xh = rand(Int32(-9):Int32(9), num_elems) - y = array_from_host(xh) - AK.accumulate!(+, y; prefer_threads, init, inclusive=false, block_size, alg) - @test Array(y) == (cumsum(xh) .- xh) .+ init - end + for _ in 1:200 + num_elems = rand(513:100_000) + block_size = rand([16, 32, 64, 128, 256]) + init = rand(Int32(-100):Int32(100)) + xh = rand(Int32(-9):Int32(9), num_elems) + y = array_from_host(xh) + AK.accumulate!(+, y; prefer_threads, init, inclusive=false, block_size, alg) + @test Array(y) == (cumsum(xh) .- xh) .+ init end # Large inclusive scan From d4a616008480ed49abd18cd3f95479b41721126b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 30 Jun 2026 16:03:29 +0200 Subject: [PATCH 2/3] fix(scan): correct DecoupledLookback exclusive carries --- src/accumulate/accumulate.jl | 5 ++-- src/accumulate/accumulate_1d_gpu.jl | 38 ++++++++++++++++++----------- test/generic/accumulate.jl | 2 +- 3 files changed, 28 insertions(+), 17 deletions(-) diff --git a/src/accumulate/accumulate.jl b/src/accumulate/accumulate.jl index e60be03..b2d10a2 100644 --- a/src/accumulate/accumulate.jl +++ b/src/accumulate/accumulate.jl @@ -105,7 +105,9 @@ The temporaries are only used for the 1D case (`dims=nothing`): `temp` stores pe `temp_flags` is only used for the `DecoupledLookback()` algorithm for flagging if blocks are ready; they should both have at least `(length(v) + 2 * block_size - 1) ÷ (2 * block_size)` elements; also, `eltype(v) === eltype(temp)` is required; the elements in `temp_flags` can be any integers, but -`UInt8` is used by default to reduce memory usage. +`UInt8` is used by default to reduce memory usage. `DecoupledLookback()` also needs cumulative +per-block aggregates; if `temp` has at least twice as many elements, its second half is used for them, +otherwise they are allocated internally. # Examples Example computing an inclusive prefix sum (the typical GPU "scan"): @@ -231,4 +233,3 @@ function accumulate( ) vcopy end - diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index 31fee7e..050c21e 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -100,10 +100,7 @@ end d = d << 0x1 end - # For exclusive ScanPrefixes scans, the local exclusive output is already correct. - # DecoupledLookback still shifts non-first blocks because _accumulate_previous! - # expects each block's last value to be globally inclusive. - if inclusive || (iblock != 0x0 && !isnothing(flags)) + if inclusive # To compute an inclusive scan, shift elements left... @synchronize() t1 = temp[ai + bank_offset_a + 0x1] @@ -132,7 +129,7 @@ end # Known at compile-time; used in the first pass of the ScanPrefixes algorithm if !isnothing(prefixes) - if isnothing(flags) && !inclusive + if !inclusive last_global = block_offset + bi if last_global < len prefixes[iblock + 0x1] = op(temp[bi + bank_offset_b + 0x1], v[last_global + 0x1]) @@ -160,7 +157,7 @@ end @kernel cpu=false inbounds=true unsafe_indices=true function _accumulate_previous!( - op, v, flags, @Const(prefixes), + op, v, flags, @Const(prefixes), aggregates, ) len = length(v) @@ -182,9 +179,8 @@ end while inspected_block >= 0x0 # Opportunistic: a previous block finished everything if UnsafeAtomics.load(pointer(flags, inspected_block + 0x1), UnsafeAtomics.monotonic) == ACC_FLAG_A - UnsafeAtomics.fence(UnsafeAtomics.acquire) # (fence before reading from v) - # Previous blocks (except last) always have filled values in v, so index is inbounds - running_prefix = op(running_prefix, v[(inspected_block + 0x1) * block_size * 0x2]) + UnsafeAtomics.fence(UnsafeAtomics.acquire) # (fence before reading from aggregates) + running_prefix = op(running_prefix, aggregates[inspected_block + 0x1]) break else running_prefix = op(running_prefix, prefixes[inspected_block + 0x1]) @@ -204,12 +200,17 @@ end v[block_offset + bi + 0x1] = op(running_prefix, v[block_offset + bi + 0x1]) end + if ithread == 0x0 + aggregates[iblock + 0x1] = op(running_prefix, prefixes[iblock + 0x1]) + end + # Set flag for "aggregate of all prefixes up to this block finished" # There are two synchronization concerns here: - # 1. Withing a group we want to ensure that all writed to `v` have occured before setting the flag. - # 2. Between groups we need to use a fence and atomic load/store to ensure that memory operations are not re-ordered + # 1. Within a group we want to ensure that all writes to `v` have occurred before setting the flag. + # 2. Between groups we need to use a fence and atomic load/store to ensure that memory operations are not re-ordered. @synchronize() # within-block - # Note: This fence is needed to ensure that the flag is not set before copying into v. + # Note: This fence is needed to ensure that the flag is not set before copying into v + # and publishing the aggregate prefix. # See https://doc.rust-lang.org/std/sync/atomic/fn.fence.html # for more details. # We use the happens-before relation between stores to `v` and the store to `flags`. @@ -295,10 +296,16 @@ function accumulate_1d_gpu!( if isnothing(temp) prefixes = similar(v, eltype(v), num_blocks) + aggregates = nothing else @argcheck eltype(temp) === eltype(v) @argcheck length(temp) >= num_blocks - prefixes = temp + prefixes = @view temp[1:num_blocks] + if length(temp) >= 2 * num_blocks + aggregates = @view temp[(num_blocks + 1):(2 * num_blocks)] + else + aggregates = nothing + end end if isnothing(temp_flags) @@ -314,8 +321,11 @@ function accumulate_1d_gpu!( ndrange=num_blocks * block_size) if num_blocks > 1 + if isnothing(aggregates) + aggregates = similar(v, eltype(v), num_blocks) + end kernel2! = _accumulate_previous!(backend, block_size) - kernel2!(op, v, flags, prefixes, + kernel2!(op, v, flags, prefixes, aggregates, ndrange=(num_blocks - 1) * block_size) end diff --git a/test/generic/accumulate.jl b/test/generic/accumulate.jl index d308f5a..05950e1 100644 --- a/test/generic/accumulate.jl +++ b/test/generic/accumulate.jl @@ -33,7 +33,7 @@ ALGS = AK.AccumulateAlgorithm[AK.ScanPrefixes()] @test all(yh .== 0:length(yh) - 1) end - # Non-uniform data exposes ScanPrefixes block-carry bugs that all-ones data masks. + # Non-uniform data exposes block-carry bugs that all-ones data masks. for _ in 1:200 num_elems = rand(513:100_000) block_size = rand([16, 32, 64, 128, 256]) From 09ff25c5a3680ea7ad689e34e6a40196bb8469ea Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 30 Jun 2026 12:00:21 -0400 Subject: [PATCH 3/3] docs(scan): document DecoupledLookback coherence bug and threadfence need MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- src/accumulate/accumulate_1d_gpu.jl | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index 050c21e..0b4ec6f 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -174,6 +174,25 @@ end block_offset = iblock * block_size * 0x2 # Processing two elements per thread # Each block looks back to find running prefix sum + # + # KNOWN BUG / TODO: this decoupled-lookback publish/consume protocol is not yet memory-coherent + # across blocks, and fails ~40% of the time on smaller GPUs (e.g. A100 MIG 1g.5gb, sm_80, ~14 + # SMs) for both inclusive and exclusive scans. Two problems: + # + # 1. Coherence. `aggregates[i]` below is read with an ordinary (L1-cacheable) load while the + # companion `flags[i]` is read atomically. A consumer can observe a fresh `flags[i] == + # ACC_FLAG_A` yet read a stale L1 copy of `aggregates[i]`. The aggregate must be accessed + # with an L1-bypassing (relaxed/monotonic atomic) load and store instead, like the flag. + # + # 2. Ordering. We need the producer's aggregate store to be globally visible *before* its flag + # store, and the consumer's aggregate load to happen *after* its flag load. That is a + # device-scope acquire/release ordering. `UnsafeAtomics.fence` (system scope) and + # acquire/release atomic load/store both FAIL TO COMPILE on recent toolchains: the NVPTX + # backend (LLVM 18, sm_80) cannot select scoped fences nor ordered atomics — only + # `monotonic` selects. The portable fix is a native device threadfence (`membar.gl` / + # `CUDA.threadfence()`), which we do not yet have as a backend-agnostic primitive in + # KernelAbstractions. It would need to be exposed here (e.g. an overridable `_decoupled_fence()` + # provided per-backend via `@device_override` in CUDA/AMDGPU package extensions). running_prefix = prefixes[iblock - 0x1 + 0x1] inspected_block = signed(typeof(iblock))(iblock) - 0x2 while inspected_block >= 0x0