From 3f0ee9a8984a51e6302446f68c5deb4be89be2b2 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Fri, 21 Aug 2026 23:12:59 +0530 Subject: [PATCH 1/6] feat(scan): overridable device fence for DecoupledLookback Replace the raw UnsafeAtomics.fence(acquire/release) in the DecoupledLookback lookback with an overridable `_decoupled_fence()` (generic no-op fallback). A plain fence is not device-scope and fails to select on recent NVPTX, so cross-block publish/consume was not coherent: correct on high-occupancy GPUs but racy where blocks run in waves. Each GPU backend supplies a native device fence via `@device_override` in its package extension. Co-Authored-By: Claude Opus 4.8 --- src/accumulate/accumulate_1d_gpu.jl | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index 4de985ed..9189cc1f 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/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), From 6ea9b01d6f6e0fa037694c5f618342e8256a55ab Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Fri, 21 Aug 2026 23:12:59 +0530 Subject: [PATCH 2/6] feat(scan): device fence extensions for all GPU backends Override `_decoupled_fence()` with a native device-scope fence per backend: CUDA threadfence; AMDGPU agent-scope seq_cst fence; oneAPI and OpenCL/POCL a SPIR-V device-scope atomic_work_item_fence; Metal an atomic_thread_fence over device memory at device scope (Metal 3.2+). Each extension is a single override so it lifts cleanly into KernelAbstractions if a portable fence lands there. Co-Authored-By: Claude Opus 4.8 --- Project.toml | 12 ++++++++++++ ext/AcceleratedKernelsAMDGPUExt.jl | 11 +++++++++++ ext/AcceleratedKernelsCUDAExt.jl | 9 +++++++++ ext/AcceleratedKernelsMetalExt.jl | 10 ++++++++++ ext/AcceleratedKernelsOpenCLExt.jl | 12 ++++++++++++ ext/AcceleratedKernelsoneAPIExt.jl | 7 +++++++ 6 files changed, 61 insertions(+) create mode 100644 ext/AcceleratedKernelsAMDGPUExt.jl create mode 100644 ext/AcceleratedKernelsCUDAExt.jl create mode 100644 ext/AcceleratedKernelsMetalExt.jl create mode 100644 ext/AcceleratedKernelsOpenCLExt.jl diff --git a/Project.toml b/Project.toml index db27fa45..07b8c214 100644 --- a/Project.toml +++ b/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" diff --git a/ext/AcceleratedKernelsAMDGPUExt.jl b/ext/AcceleratedKernelsAMDGPUExt.jl new file mode 100644 index 00000000..f5a5cb40 --- /dev/null +++ b/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 diff --git a/ext/AcceleratedKernelsCUDAExt.jl b/ext/AcceleratedKernelsCUDAExt.jl new file mode 100644 index 00000000..daa6b7c6 --- /dev/null +++ b/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 diff --git a/ext/AcceleratedKernelsMetalExt.jl b/ext/AcceleratedKernelsMetalExt.jl new file mode 100644 index 00000000..0e3f8f94 --- /dev/null +++ b/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 diff --git a/ext/AcceleratedKernelsOpenCLExt.jl b/ext/AcceleratedKernelsOpenCLExt.jl new file mode 100644 index 00000000..870d0a87 --- /dev/null +++ b/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 diff --git a/ext/AcceleratedKernelsoneAPIExt.jl b/ext/AcceleratedKernelsoneAPIExt.jl index 1746ca7e..e34db7d5 100644 --- a/ext/AcceleratedKernelsoneAPIExt.jl +++ b/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( From 10661b10185ccb6e653e1ce8f51b42e7d2090e27 Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Fri, 21 Aug 2026 23:57:15 +0530 Subject: [PATCH 3/6] test(scan): many-block DecoupledLookback coherence stress Add a small-tile (block_size 16-64, one item per thread) non-uniform scan loop in both directions, run for every algorithm. It maximises the number of inter-block publish/consume handoffs so DecoupledLookback exercises the device fence on many blocks; a fence that is not device scoped drops whole-block carries and fails here. Co-Authored-By: Claude Opus 4.8 --- test/generic/accumulate.jl | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/test/generic/accumulate.jl b/test/generic/accumulate.jl index 7ed00893..cadf4684 100644 --- a/test/generic/accumulate.jl +++ b/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 From 5cf3169401dc86b15dc76e9de014cd6e3142e19c Mon Sep 17 00:00:00 2001 From: shreyas-omkar Date: Tue, 25 Aug 2026 16:32:35 +0530 Subject: [PATCH 4/6] Address review: empty _decoupled_fence, fix AMDGPU and Metal compat bounds --- Project.toml | 4 ++-- src/accumulate/accumulate_1d_gpu.jl | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 07b8c214..c71289d3 100644 --- a/Project.toml +++ b/Project.toml @@ -27,14 +27,14 @@ AcceleratedKernelsOpenCLExt = "OpenCL" AcceleratedKernelsoneAPIExt = "oneAPI" [compat] -AMDGPU = "1, 2" +AMDGPU = "1.3.4, 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" +Metal = "1.10.1" OpenCL = "0.10" SIMD = "3" UnsafeAtomics = "0.3.0" diff --git a/src/accumulate/accumulate_1d_gpu.jl b/src/accumulate/accumulate_1d_gpu.jl index 9189cc1f..f4d1a481 100644 --- a/src/accumulate/accumulate_1d_gpu.jl +++ b/src/accumulate/accumulate_1d_gpu.jl @@ -12,7 +12,7 @@ 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 +function _decoupled_fence end # Register-raking block scan with striped loads and stores. From b4ac428622238537cd2b5d72b3e3c9b725bc358e Mon Sep 17 00:00:00 2001 From: Christian <28689358+christiangnrd@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:41:06 -0300 Subject: [PATCH 5/6] Extend CUDACore not CUDA --- Project.toml | 6 +++--- ext/AcceleratedKernelsCUDACoreExt.jl | 9 +++++++++ ext/AcceleratedKernelsCUDAExt.jl | 9 --------- 3 files changed, 12 insertions(+), 12 deletions(-) create mode 100644 ext/AcceleratedKernelsCUDACoreExt.jl delete mode 100644 ext/AcceleratedKernelsCUDAExt.jl diff --git a/Project.toml b/Project.toml index c71289d3..6e7b8a15 100644 --- a/Project.toml +++ b/Project.toml @@ -14,14 +14,14 @@ UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f" [weakdeps] AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e" -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" +CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" Metal = "dde4c033-4e86-420c-a63e-0dd931031962" OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] AcceleratedKernelsAMDGPUExt = "AMDGPU" -AcceleratedKernelsCUDAExt = "CUDA" +AcceleratedKernelsCUDACoreExt = "CUDACore" AcceleratedKernelsMetalExt = "Metal" AcceleratedKernelsOpenCLExt = "OpenCL" AcceleratedKernelsoneAPIExt = "oneAPI" @@ -30,7 +30,7 @@ AcceleratedKernelsoneAPIExt = "oneAPI" AMDGPU = "1.3.4, 2" ArgCheck = "2" Atomix = "0.1, 1" -CUDA = "5, 6" +CUDACore = "6" GPUArraysCore = "0.2.0" KernelAbstractions = "0.9.34, 0.10" Markdown = "1" diff --git a/ext/AcceleratedKernelsCUDACoreExt.jl b/ext/AcceleratedKernelsCUDACoreExt.jl new file mode 100644 index 00000000..1b5cbe64 --- /dev/null +++ b/ext/AcceleratedKernelsCUDACoreExt.jl @@ -0,0 +1,9 @@ +module AcceleratedKernelsCUDACoreExt + +using CUDACore +import AcceleratedKernels as AK + +# Device-scope fence for the DecoupledLookback scan. +CUDACore.@device_override AK._decoupled_fence() = CUDACore.threadfence() + +end diff --git a/ext/AcceleratedKernelsCUDAExt.jl b/ext/AcceleratedKernelsCUDAExt.jl deleted file mode 100644 index daa6b7c6..00000000 --- a/ext/AcceleratedKernelsCUDAExt.jl +++ /dev/null @@ -1,9 +0,0 @@ -module AcceleratedKernelsCUDAExt - -using CUDA -import AcceleratedKernels as AK - -# Device-scope fence for the DecoupledLookback scan. -CUDA.@device_override AK._decoupled_fence() = CUDA.threadfence() - -end From f3406d24128b46e8d50086f5bad3723aa62d919d Mon Sep 17 00:00:00 2001 From: Christian <28689358+christiangnrd@users.noreply.github.com> Date: Tue, 25 Aug 2026 09:45:37 -0300 Subject: [PATCH 6/6] Rename extensions to follow current recommendations --- Project.toml | 10 +++++----- ext/{AcceleratedKernelsAMDGPUExt.jl => AMDGPUExt.jl} | 2 +- ...AcceleratedKernelsCUDACoreExt.jl => CUDACoreExt.jl} | 2 +- ext/{AcceleratedKernelsMetalExt.jl => MetalExt.jl} | 2 +- ext/{AcceleratedKernelsOpenCLExt.jl => OpenCLExt.jl} | 2 +- ext/{AcceleratedKernelsoneAPIExt.jl => oneAPIExt.jl} | 4 ++-- 6 files changed, 11 insertions(+), 11 deletions(-) rename ext/{AcceleratedKernelsAMDGPUExt.jl => AMDGPUExt.jl} (88%) rename ext/{AcceleratedKernelsCUDACoreExt.jl => CUDACoreExt.jl} (82%) rename ext/{AcceleratedKernelsMetalExt.jl => MetalExt.jl} (88%) rename ext/{AcceleratedKernelsOpenCLExt.jl => OpenCLExt.jl} (91%) rename ext/{AcceleratedKernelsoneAPIExt.jl => oneAPIExt.jl} (92%) diff --git a/Project.toml b/Project.toml index 6e7b8a15..5bf1cc2e 100644 --- a/Project.toml +++ b/Project.toml @@ -20,11 +20,11 @@ OpenCL = "08131aa3-fb12-5dee-8b74-c09406e224a2" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" [extensions] -AcceleratedKernelsAMDGPUExt = "AMDGPU" -AcceleratedKernelsCUDACoreExt = "CUDACore" -AcceleratedKernelsMetalExt = "Metal" -AcceleratedKernelsOpenCLExt = "OpenCL" -AcceleratedKernelsoneAPIExt = "oneAPI" +AMDGPUExt = "AMDGPU" +CUDACoreExt = "CUDACore" +MetalExt = "Metal" +OpenCLExt = "OpenCL" +oneAPIExt = "oneAPI" [compat] AMDGPU = "1.3.4, 2" diff --git a/ext/AcceleratedKernelsAMDGPUExt.jl b/ext/AMDGPUExt.jl similarity index 88% rename from ext/AcceleratedKernelsAMDGPUExt.jl rename to ext/AMDGPUExt.jl index f5a5cb40..92758ef2 100644 --- a/ext/AcceleratedKernelsAMDGPUExt.jl +++ b/ext/AMDGPUExt.jl @@ -1,4 +1,4 @@ -module AcceleratedKernelsAMDGPUExt +module AMDGPUExt using AMDGPU import UnsafeAtomics diff --git a/ext/AcceleratedKernelsCUDACoreExt.jl b/ext/CUDACoreExt.jl similarity index 82% rename from ext/AcceleratedKernelsCUDACoreExt.jl rename to ext/CUDACoreExt.jl index 1b5cbe64..9f0a3c25 100644 --- a/ext/AcceleratedKernelsCUDACoreExt.jl +++ b/ext/CUDACoreExt.jl @@ -1,4 +1,4 @@ -module AcceleratedKernelsCUDACoreExt +module CUDACoreExt using CUDACore import AcceleratedKernels as AK diff --git a/ext/AcceleratedKernelsMetalExt.jl b/ext/MetalExt.jl similarity index 88% rename from ext/AcceleratedKernelsMetalExt.jl rename to ext/MetalExt.jl index 0e3f8f94..d5aeee29 100644 --- a/ext/AcceleratedKernelsMetalExt.jl +++ b/ext/MetalExt.jl @@ -1,4 +1,4 @@ -module AcceleratedKernelsMetalExt +module MetalExt using Metal import AcceleratedKernels as AK diff --git a/ext/AcceleratedKernelsOpenCLExt.jl b/ext/OpenCLExt.jl similarity index 91% rename from ext/AcceleratedKernelsOpenCLExt.jl rename to ext/OpenCLExt.jl index 870d0a87..961b26f0 100644 --- a/ext/AcceleratedKernelsOpenCLExt.jl +++ b/ext/OpenCLExt.jl @@ -1,4 +1,4 @@ -module AcceleratedKernelsOpenCLExt +module OpenCLExt using OpenCL using OpenCL: method_table # used by OpenCL.@device_override diff --git a/ext/AcceleratedKernelsoneAPIExt.jl b/ext/oneAPIExt.jl similarity index 92% rename from ext/AcceleratedKernelsoneAPIExt.jl rename to ext/oneAPIExt.jl index e34db7d5..66c2cea0 100644 --- a/ext/AcceleratedKernelsoneAPIExt.jl +++ b/ext/oneAPIExt.jl @@ -1,4 +1,4 @@ -module AcceleratedKernelsoneAPIExt +module oneAPIExt using oneAPI @@ -44,4 +44,4 @@ function AK.all( end -end # module AcceleratedKernelsoneAPIExt \ No newline at end of file +end # module oneAPIExt