diff --git a/Project.toml b/Project.toml index b5d60380..be5aea0a 100644 --- a/Project.toml +++ b/Project.toml @@ -11,13 +11,16 @@ Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a" UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f" [weakdeps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] +AcceleratedKernelsCUDAExt = "CUDA" AcceleratedKernelsoneAPIExt = "oneAPI" [compat] ArgCheck = "2" +CUDA = "5, 6" GPUArraysCore = "0.2.0" KernelAbstractions = "0.9.34, 0.10" Markdown = "1" diff --git a/ext/AcceleratedKernelsCUDAExt.jl b/ext/AcceleratedKernelsCUDAExt.jl new file mode 100644 index 00000000..29f33fdc --- /dev/null +++ b/ext/AcceleratedKernelsCUDAExt.jl @@ -0,0 +1,16 @@ +module AcceleratedKernelsCUDAExt + + +import AcceleratedKernels as AK +import CUDA +using CUDA: @device_override + + +# On CUDA the NVPTX backend does not select scoped atomic fences, so `UnsafeAtomics.fence` +# with acquire/release ordering (the generic `AK._decoupled_fence`) either fails to lower or +# is not device-scope coherent. Provide the decoupled-lookback device fence with the native +# `membar.gl` threadfence instead. +@device_override @inline AK._decoupled_fence() = CUDA.threadfence() + + +end # module AcceleratedKernelsCUDAExt diff --git a/src/accumulate/accumulate.jl b/src/accumulate/accumulate.jl index e60be03f..b2d10a2a 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 31fee7ea..7bdf8648 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -1,8 +1,11 @@ const ACC_NUM_BANKS::UInt8 = 32 const ACC_LOG_NUM_BANKS::UInt8 = 5 -const ACC_FLAG_A::UInt8 = 0 # Aggregate of all previous prefixes finished -const ACC_FLAG_P::UInt8 = 1 # Only current block's prefix available +# 32-bit, not 8-bit: the decoupled-lookback flag is accessed with an atomic load/store, and +# SPIR-V / oneAPI Level Zero has no 8-bit atomics (a UInt8 flag array fails the module build +# with ZE_RESULT_ERROR_MODULE_BUILD_FAILURE). UInt32 atomics are supported everywhere. +const ACC_FLAG_A::UInt32 = 0 # Aggregate of all previous prefixes finished +const ACC_FLAG_P::UInt32 = 1 # Only current block's prefix available @inline function conflict_free_offset(n) @@ -12,7 +15,17 @@ const ACC_FLAG_P::UInt8 = 1 # Only current block's prefix available end -@kernel cpu=false inbounds=true unsafe_indices=true function _accumulate_block!( +# NOTE: this kernel is declared `inbounds=false` (bounds-checked) rather than the usual +# `inbounds=true`. On AMDGPU/ROCm (observed on RDNA4 gfx1200), forcing `@inbounds` over the +# whole body miscompiles the multi-block path into an illegal-address access for any input +# spanning more than one block (n > 2*block_size), faulting both ScanPrefixes and +# DecoupledLookback. It is a codegen bug, not a real out-of-bounds — the same kernel is +# correct under `--check-bounds=yes`, which does not throw. So we bounds-check by default and +# re-apply `@inbounds` only to the hot, provably-in-bounds accesses (the Blelloch tree-scan +# over shared memory and the guarded coalesced global load/store); the cold per-block carry +# region below is left checked, which is where the miscompile lives and costs ~nothing since +# only one thread per block runs it. +@kernel cpu=false inbounds=false unsafe_indices=true function _accumulate_block!( op, v, init, neutral, inclusive, flags, prefixes, # one per block @@ -42,13 +55,13 @@ end bank_offset_a = conflict_free_offset(ai) bank_offset_b = conflict_free_offset(bi) - if block_offset + ai < len + @inbounds if block_offset + ai < len temp[ai + bank_offset_a + 0x1] = v[block_offset + ai + 0x1] else temp[ai + bank_offset_a + 0x1] = neutral end - if block_offset + bi < len + @inbounds if block_offset + bi < len temp[bi + bank_offset_b + 0x1] = v[block_offset + bi + 0x1] else temp[bi + bank_offset_b + 0x1] = neutral @@ -67,7 +80,7 @@ end _ai += conflict_free_offset(_ai) _bi += conflict_free_offset(_bi) - temp[_bi + 0x1] = op(temp[_bi + 0x1], temp[_ai + 0x1]) + @inbounds temp[_bi + 0x1] = op(temp[_bi + 0x1], temp[_ai + 0x1]) end offset = offset << 0x1 @@ -92,18 +105,17 @@ end _ai += conflict_free_offset(_ai) _bi += conflict_free_offset(_bi) - t = temp[_ai + 0x1] - temp[_ai + 0x1] = temp[_bi + 0x1] - temp[_bi + 0x1] = op(temp[_bi + 0x1], t) + @inbounds begin + t = temp[_ai + 0x1] + temp[_ai + 0x1] = temp[_bi + 0x1] + temp[_bi + 0x1] = op(temp[_bi + 0x1], t) + end 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 +144,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]) @@ -150,17 +162,76 @@ end end end - if block_offset + ai < len + @inbounds if block_offset + ai < len v[block_offset + ai + 0x1] = temp[ai + bank_offset_a + 0x1] end - if block_offset + bi < len + @inbounds if block_offset + bi < len v[block_offset + bi + 0x1] = temp[bi + bank_offset_b + 0x1] end end +# ─── Decoupled-lookback cross-block coherence primitives ───────────────────── +# +# The decoupled-lookback protocol has each block publish its inclusive aggregate plus a +# companion flag, which later blocks consume. For this to be correct across blocks on a +# GPU, two things are required: +# +# 1. Coherence. The aggregate must be read/written with an L1-bypassing, device-scope +# relaxed (`monotonic`) atomic load/store, exactly like the flag. A plain cached load +# can return a stale aggregate even after its flag is observed as published. +# 2. Ordering. A device-scope memory fence must separate the aggregate store from the +# flag store (producer), and the flag load from the aggregate load (consumer). +# +# `_decoupled_fence()` is that device-scope fence. The generic definition uses an +# acquire-release atomic fence, which lowers on OpenCL/SPIR-V, Metal, oneAPI and AMDGPU. +# On CUDA the NVPTX backend does not select scoped atomic fences (only `monotonic` does), +# so the CUDA package extension overrides this with `CUDA.threadfence()` (`membar.gl`). +@inline _decoupled_fence() = UnsafeAtomics.fence(UnsafeAtomics.acq_rel) + +# Unsigned integer type with the same width as T, or `Nothing` if T is not a 1/2/4/8-byte +# bitstype — in which case we fall back to a plain (non-atomic) access. +@inline _relaxed_utype(::Type{T}) where {T} = + sizeof(T) == 1 ? UInt8 : + sizeof(T) == 2 ? UInt16 : + sizeof(T) == 4 ? UInt32 : + sizeof(T) == 8 ? UInt64 : Nothing + +# Reinterpret a device pointer to a different element type while PRESERVING its address +# space. `reinterpret(Ptr{U}, pointer(a, i))` drops the pointer to the generic address space +# (AS 0), which NVPTX/AMDGPU tolerate but the SPIR-V translator (oneAPI Level Zero) rejects +# ("Failed to translate LLVM code to SPIR-V"). Keeping the LLVMPtr's own address space is +# legal everywhere. +@inline _as_llvmptr(::Type{U}, p::Core.LLVMPtr{T, AS}) where {U, T, AS} = + reinterpret(Core.LLVMPtr{U, AS}, p) + +# Device-scope relaxed (L1-bypassing) atomic load/store of a[i], reinterpreting through a +# same-width unsigned integer so float element types are supported. The `U === Nothing` +# branch is resolved at compile time. +@inline function _load_relaxed(a, i) + U = _relaxed_utype(eltype(a)) + if U === Nothing + return a[i] + else + p = _as_llvmptr(U, pointer(a, i)) + return reinterpret(eltype(a), UnsafeAtomics.load(p, UnsafeAtomics.monotonic)) + end +end + +@inline function _store_relaxed!(a, i, x::T) where {T} + U = _relaxed_utype(T) + if U === Nothing + a[i] = x + else + p = _as_llvmptr(U, pointer(a, i)) + UnsafeAtomics.store!(p, reinterpret(U, x), UnsafeAtomics.monotonic) + end + return nothing +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) @@ -176,15 +247,21 @@ end ithread = @index(Local, Linear) - 0x1 block_offset = iblock * block_size * 0x2 # Processing two elements per thread - # Each block looks back to find running prefix sum + # Each block looks back over previous blocks to find its running prefix. A previous + # block that has published its full aggregate (flag == ACC_FLAG_A) lets us stop early; + # otherwise we fall back to that block's own prefix and keep looking further back. + # + # Cross-block coherence (see `_decoupled_fence` above): the flag is read with a relaxed + # atomic; once ACC_FLAG_A is observed, a device-scope fence orders the subsequent + # aggregate load *after* the flag load, and the aggregate itself is read with a relaxed + # (L1-bypassing) atomic so it cannot be a stale cached value. running_prefix = prefixes[iblock - 0x1 + 0x1] inspected_block = signed(typeof(iblock))(iblock) - 0x2 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]) + _decoupled_fence() # order the aggregate load after the flag load + running_prefix = op(running_prefix, _load_relaxed(aggregates, inspected_block + 0x1)) break else running_prefix = op(running_prefix, prefixes[inspected_block + 0x1]) @@ -204,16 +281,22 @@ end v[block_offset + bi + 0x1] = op(running_prefix, v[block_offset + bi + 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 + # Publish this block's inclusive aggregate with a relaxed atomic store, so consumers + # never read a stale L1 copy once they see our flag. + if ithread == 0x0 + _store_relaxed!(aggregates, iblock + 0x1, op(running_prefix, prefixes[iblock + 0x1])) + end + + # Set flag for "aggregate of all prefixes up to this block finished". + # Two synchronization concerns: + # 1. Within a group, all writes to `v` must have occurred before setting the flag. + # 2. Between groups, the aggregate store must be globally visible before the flag store, + # so a consumer observing ACC_FLAG_A is guaranteed to see the published aggregate. @synchronize() # within-block - # Note: This fence is needed to ensure that the flag is not set before copying into v. - # 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`. - UnsafeAtomics.fence(UnsafeAtomics.release) + # Device-scope fence: orders the writes to `v` and the aggregate store *before* the flag + # store. Together with the consumer-side fence this gives a happens-before relation + # between publishing the aggregate and observing the flag. + _decoupled_fence() if ithread == 0x0 UnsafeAtomics.store!(pointer(flags, iblock + 0x1), convert(eltype(flags), ACC_FLAG_A), UnsafeAtomics.monotonic) end @@ -298,11 +381,16 @@ function accumulate_1d_gpu!( else @argcheck eltype(temp) === eltype(v) @argcheck length(temp) >= num_blocks - prefixes = temp + prefixes = @view temp[1:num_blocks] end + # `aggregates` is always its own contiguous array (never a view): the decoupled-lookback + # kernel takes `pointer(aggregates, i)` for relaxed atomic access, which is only valid on + # a plain device array. + aggregates = nothing + if isnothing(temp_flags) - flags = similar(v, UInt8, num_blocks) + flags = similar(v, UInt32, num_blocks) # UInt32 (not UInt8): see ACC_FLAG_A note — SPIR-V has no 8-bit atomics else @argcheck eltype(temp_flags) <: Integer @argcheck length(temp_flags) >= num_blocks @@ -314,8 +402,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 d1c2adf1..6463adb8 100644 --- a/test/generic/accumulate.jl +++ b/test/generic/accumulate.jl @@ -33,20 +33,15 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback()) @test all(yh .== 0:length(yh) - 1) 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 + # 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]) + 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