Fix DecoupledLookback with a device-scope memory fence

Project.toml

@@ -13,17 +13,29 @@ SIMD = "fdea26ae-647d-5447-a871-4b548cad5224"
 UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 
 [weakdeps]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+Metal = "dde4c033-4e86-420c-a63e-0dd931031962"
+OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2"
 oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 
 [extensions]
+AcceleratedKernelsAMDGPUExt = "AMDGPU"
+AcceleratedKernelsCUDAExt = "CUDA"
+AcceleratedKernelsMetalExt = "Metal"
+AcceleratedKernelsOpenCLExt = "OpenCL"
 AcceleratedKernelsoneAPIExt = "oneAPI"
 
 [compat]
+AMDGPU = "1, 2"
 ArgCheck = "2"
 Atomix = "0.1, 1"
+CUDA = "5, 6"
 GPUArraysCore = "0.2.0"
 KernelAbstractions = "0.9.34, 0.10"
 Markdown = "1"
+Metal = "1.10, 2"
+OpenCL = "0.10"
 SIMD = "3"
 UnsafeAtomics = "0.3.0"
 julia = "1.10"

ext/AcceleratedKernelsAMDGPUExt.jl

@@ -0,0 +1,11 @@
+module AcceleratedKernelsAMDGPUExt
+
+using AMDGPU
+import UnsafeAtomics
+import AcceleratedKernels as AK
+
+# Device-scope (agent) fence for the DecoupledLookback scan.
+AMDGPU.Device.@device_override AK._decoupled_fence() =
+    UnsafeAtomics.fence(UnsafeAtomics.seq_cst, AMDGPU.syncscope_agent)
+
+end

ext/AcceleratedKernelsCUDAExt.jl

@@ -0,0 +1,9 @@
+module AcceleratedKernelsCUDAExt
+
+using CUDA
+import AcceleratedKernels as AK
+
+# Device-scope fence for the DecoupledLookback scan.
+CUDA.@device_override AK._decoupled_fence() = CUDA.threadfence()
+
+end

ext/AcceleratedKernelsMetalExt.jl

@@ -0,0 +1,10 @@
+module AcceleratedKernelsMetalExt
+
+using Metal
+import AcceleratedKernels as AK
+
+# Device-scope fence for the DecoupledLookback scan (Metal 3.2+).
+Metal.@device_override AK._decoupled_fence() =
+    Metal.atomic_thread_fence(Metal.MemoryFlagDevice, Metal.memory_order_seq_cst, Metal.thread_scope_device)
+
+end

ext/AcceleratedKernelsOpenCLExt.jl

@@ -0,0 +1,12 @@
+module AcceleratedKernelsOpenCLExt
+
+using OpenCL
+using OpenCL: method_table          # used by OpenCL.@device_override
+import AcceleratedKernels as AK
+
+# Device-scope SPIR-V fence for the DecoupledLookback scan (also the POCL path).
+const SPIRV = OpenCL.SPIRVIntrinsics
+OpenCL.@device_override AK._decoupled_fence() =
+    SPIRV.atomic_work_item_fence(SPIRV.GLOBAL_MEM_FENCE, SPIRV.memory_order_seq_cst, SPIRV.memory_scope_device)
+
+end

ext/AcceleratedKernelsoneAPIExt.jl

@@ -2,9 +2,16 @@ module AcceleratedKernelsoneAPIExt
 
 
 using oneAPI
+using oneAPI: method_table          # used by oneAPI.@device_override
 import AcceleratedKernels as AK
 
 
+# Device-scope SPIR-V fence for the DecoupledLookback scan.
+const SPIRV = oneAPI.SPIRVIntrinsics
+oneAPI.@device_override AK._decoupled_fence() =
+    SPIRV.atomic_work_item_fence(SPIRV.GLOBAL_MEM_FENCE, SPIRV.memory_order_seq_cst, SPIRV.memory_scope_device)
+
+
 # On oneAPI, use the MapReduce algorithm by default as on some Intel GPUs ConcurrentWrite hangs
 # the device.
 function AK.any(

src/accumulate/accumulate_1d_gpu.jl

@@ -10,6 +10,11 @@ const ACC_FLAG_P::UInt8 = 1             # Only current block's prefix available
 end
 
 
+# Device-scope memory fence for the DecoupledLookback scan. Each GPU backend overrides it with a
+# native device fence in its package extension; a plain UnsafeAtomics.fence is not device scoped.
+@inline _decoupled_fence() = nothing
+
+
 # Register-raking block scan with striped loads and stores.
 @kernel cpu=false inbounds=true unsafe_indices=true function _accumulate_block!(
     op, v, init, neutral,
@@ -147,7 +152,7 @@ end
             UnsafeAtomics.monotonic,
         )
         if flag == ACC_FLAG_A
-            UnsafeAtomics.fence(UnsafeAtomics.acquire)
+            _decoupled_fence()          # acquire: order the `v` read after the flag load
             running_prefix = op(running_prefix, v[(inspected_block + 0x1) * block_size * ITEMS])
             break
         else
@@ -169,7 +174,7 @@ end
 
     # Publish writes to `v` before marking the block complete.
     @synchronize()
-    UnsafeAtomics.fence(UnsafeAtomics.release)
+    _decoupled_fence()                  # release: order the flag store after the `v` writes
     if ithread == 0x0
         UnsafeAtomics.store!(
             pointer(flags, iblock + 0x1),

test/generic/accumulate.jl

@@ -122,6 +122,27 @@ TEST_DL && push!(ALGS, AK.DecoupledLookback())
                       temp=array_from_host(zeros(Int64, 1000)),
                       temp_flags=array_from_host(zeros(Int8, 1000)))
     @test Array(y) == 0:999
+
+    # Cross-block coherence: small tiles (block_size 16-64, 1 item/thread) maximise the number of
+    # inter-block publish/consume handoffs. For DecoupledLookback each handoff relies on the
+    # device-scope fence, so many-block non-uniform scans in both directions guard against a fence
+    # that is not device scoped (the incoherent lookback would drop whole-block carries).
+    for _ in 1:100
+        num_elems = rand(5_000:200_000)
+        block_size = rand((16, 32, 64))
+        xh = rand(Int32(-9):Int32(9), num_elems)
+
+        yi = array_from_host(xh)
+        AK.accumulate!(+, yi; prefer_threads, init=Int32(0), inclusive=true,
+                       block_size, items_per_thread=1, alg)
+        @test Array(yi) == cumsum(xh)
+
+        init = rand(Int32(-50):Int32(50))
+        ye = array_from_host(xh)
+        AK.accumulate!(+, ye; prefer_threads, init, inclusive=false,
+                       block_size, items_per_thread=1, alg)
+        @test Array(ye) == (cumsum(xh) .- xh) .+ init
+    end
 end