Add Atomix and UnsafeAtomics atomics tests

src/pocl/compiler/compilation.jl

@@ -142,6 +142,50 @@ end
 
 ## compiler implementation (configure, compile, and link)
 
+"""
+    supports_fp_atomics(caps::UInt64, ops::UInt64)
+
+Whether the `cl_ext_float_atomics` capability bitfield `caps` (see
+`dev.single_fp_atomic_capabilities` and friends) natively supports all of `ops`.
+Kernels perform atomics on both global and local memory, so callers should
+require both the `GLOBAL` and `LOCAL` bit of an operation.
+"""
+supports_fp_atomics(caps::UInt64, ops::UInt64) = caps & ops == ops
+
+const fp_atomic_add = cl.CL_DEVICE_GLOBAL_FP_ATOMIC_ADD_EXT | cl.CL_DEVICE_LOCAL_FP_ATOMIC_ADD_EXT
+const fp_atomic_min_max = cl.CL_DEVICE_GLOBAL_FP_ATOMIC_MIN_MAX_EXT | cl.CL_DEVICE_LOCAL_FP_ATOMIC_MIN_MAX_EXT
+
+"""
+    default_spirv_extensions(dev)
+
+SPIR-V extensions to permit for `dev`, as the `+`-prefixed, comma-separated string
+`SPIRVCompilerTarget` passes on to the backend via `-spirv-ext`.
+
+Listing an extension only *permits* it: nothing is emitted unless a module actually
+needs the instructions it guards, so this costs nothing for kernels that don't.
+"""
+function default_spirv_extensions(dev)
+    exts = String[]
+
+    # Floating-point atomics. Atomix/UnsafeAtomics lower `@atomic A[i] += x` and
+    # `@atomic max(A[i], x)` on floats to LLVM `atomicrmw fadd`/`fmin`/`fmax`, which the
+    # SPIR-V backend only translates when the corresponding extension is permitted:
+    #   LLVM ERROR: The atomic float instruction requires the following SPIR-V
+    #   extension: SPV_EXT_shader_atomic_float_add
+    # Enzyme's reverse mode hits this too, as it accumulates gradients with atomic fadd.
+    # The device reports native support per precision through cl_ext_float_atomics.
+    fp32 = dev.single_fp_atomic_capabilities
+    fp64 = dev.double_fp_atomic_capabilities
+    if supports_fp_atomics(fp32, fp_atomic_add) || supports_fp_atomics(fp64, fp_atomic_add)
+        push!(exts, "+SPV_EXT_shader_atomic_float_add")
+    end
+    if supports_fp_atomics(fp32, fp_atomic_min_max) || supports_fp_atomics(fp64, fp_atomic_min_max)
+        push!(exts, "+SPV_EXT_shader_atomic_float_min_max")
+    end
+
+    return join(exts, ",")
+end
+
 # cache of compiler configurations, per device (but additionally configurable via kwargs)
 const _toolchain = Ref{Any}()
 const _compiler_configs = Dict{UInt, OpenCLCompilerConfig}()
@@ -154,16 +198,24 @@ function compiler_config(dev::cl.Device; kwargs...)
     end
     return config
 end
-@noinline function _compiler_config(dev; kernel = true, name = nothing, always_inline = false, sub_group_size::Union{Nothing, Int} = 32, kwargs...)
+@noinline function _compiler_config(
+        dev; kernel = true, name = nothing, always_inline = false,
+        sub_group_size::Union{Nothing, Int} = 32,
+        extensions::Union{Nothing, String} = nothing, kwargs...
+    )
     supports_fp16 = "cl_khr_fp16" in dev.extensions
     supports_fp64 = "cl_khr_fp64" in dev.extensions
 
     if sub_group_size !== nothing && sub_group_size ∉ dev.sub_group_sizes
         error("$sub_group_size is not a valid sub-group size for this device.")
     end
 
+    if extensions === nothing
+        extensions = default_spirv_extensions(dev)
+    end
+
     # create GPUCompiler objects
-    target = SPIRVCompilerTarget(; supports_fp16, supports_fp64, validate = true, kwargs...)
+    target = SPIRVCompilerTarget(; supports_fp16, supports_fp64, extensions, validate = true, kwargs...)
     params = OpenCLCompilerParams(; sub_group_size)
     return CompilerConfig(target, params; kernel, name, always_inline)
 end

src/pocl/nanoOpenCL.jl

@@ -299,6 +299,25 @@ const CL_DEVICE_NON_UNIFORM_WORK_GROUP_SUPPORT = 0x1065
 
 const CL_DEVICE_OPENCL_C_ALL_VERSIONS = 0x1066
 
+# cl_ext_float_atomics
+const CL_DEVICE_SINGLE_FP_ATOMIC_CAPABILITIES_EXT = 0x4231
+
+const CL_DEVICE_DOUBLE_FP_ATOMIC_CAPABILITIES_EXT = 0x4232
+
+const CL_DEVICE_HALF_FP_ATOMIC_CAPABILITIES_EXT = 0x4233
+
+const CL_DEVICE_GLOBAL_FP_ATOMIC_LOAD_STORE_EXT = UInt64(1) << 0
+
+const CL_DEVICE_GLOBAL_FP_ATOMIC_ADD_EXT = UInt64(1) << 1
+
+const CL_DEVICE_GLOBAL_FP_ATOMIC_MIN_MAX_EXT = UInt64(1) << 2
+
+const CL_DEVICE_LOCAL_FP_ATOMIC_LOAD_STORE_EXT = UInt64(1) << 16
+
+const CL_DEVICE_LOCAL_FP_ATOMIC_ADD_EXT = UInt64(1) << 17
+
+const CL_DEVICE_LOCAL_FP_ATOMIC_MIN_MAX_EXT = UInt64(1) << 18
+
 const CL_DEVICE_PREFERRED_WORK_GROUP_SIZE_MULTIPLE = 0x1067
 
 const CL_DEVICE_WORK_GROUP_COLLECTIVE_FUNCTIONS_SUPPORT = 0x1068
@@ -924,6 +943,18 @@ devices(p::Platform) = devices(p, CL_DEVICE_TYPE_ALL)
         return String[string(s) for s in split(bs)]
     end
 
+    # cl_ext_float_atomics: per-precision bitfields of natively supported floating-point
+    # atomic operations (zero when the device does not expose the extension)
+    if s == :single_fp_atomic_capabilities || s == :double_fp_atomic_capabilities || s == :half_fp_atomic_capabilities
+        "cl_ext_float_atomics" in d.extensions || return zero(UInt64)
+        prop = s == :single_fp_atomic_capabilities ? CL_DEVICE_SINGLE_FP_ATOMIC_CAPABILITIES_EXT :
+            s == :double_fp_atomic_capabilities ? CL_DEVICE_DOUBLE_FP_ATOMIC_CAPABILITIES_EXT :
+            CL_DEVICE_HALF_FP_ATOMIC_CAPABILITIES_EXT
+        caps = Ref{UInt64}(0)
+        clGetDeviceInfo(d, prop, sizeof(UInt64), caps, C_NULL)
+        return caps[]
+    end
+
     if s == :platform
         result = Ref{cl_platform_id}()
         clGetDeviceInfo(d, CL_DEVICE_PLATFORM, sizeof(cl_platform_id), result, C_NULL)

test/Project.toml

@@ -10,6 +10,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 pocl_standalone_jll = "54f56a70-6062-5590-a942-1226658f6c83"
 
 [sources]

test/atomics.jl

@@ -0,0 +1,261 @@
+using KernelAbstractions
+using KernelAbstractions: @atomic, @atomicswap, @atomicreplace
+using Test
+
+import UnsafeAtomics
+
+# Atomix based kernels
+
+@kernel function atomix_add!(hist)
+    i = @index(Global, Linear)
+    j = (i - 1) % length(hist) + 1
+    @inbounds @atomic hist[j] += one(eltype(hist))
+end
+
+@kernel function atomix_max!(A)
+    i = @index(Global, Linear)
+    @inbounds @atomic max(A[1], eltype(A)(i))
+end
+
+@kernel function atomix_min!(A)
+    i = @index(Global, Linear)
+    @inbounds @atomic min(A[1], eltype(A)(i))
+end
+
+@kernel function atomix_load_store!(A, B)
+    i = @index(Global, Linear)
+    @inbounds begin
+        v = @atomic B[i]
+        @atomic A[i] = v
+    end
+end
+
+@kernel function atomix_swap!(A, B)
+    i = @index(Global, Linear)
+    @inbounds B[i] = @atomicswap A[i] = eltype(A)(i)
+end
+
+@kernel function atomix_replace!(A, success)
+    i = @index(Global, Linear)
+    T = eltype(A)
+    @inbounds begin
+        # CAS that must succeed, followed by one that must fail
+        (_, ok1) = @atomicreplace A[i] zero(T) => T(i)
+        (_, ok2) = @atomicreplace A[i] zero(T) => T(-1)
+        success[i] = ok1 & !ok2
+    end
+end
+
+@kernel function atomix_ordered!(A, B)
+    i = @index(Global, Linear)
+    T = eltype(A)
+    @inbounds begin
+        @atomic :release A[i] = T(i)
+        v = @atomic :acquire A[i]
+        @atomic :monotonic A[i] += one(T)
+        @atomic :acquire_release A[i] += one(T)
+        @atomic :sequentially_consistent A[i] += one(T)
+        B[i] = @atomicswap :acquire_release A[i] = v + T(3)
+    end
+end
+
+# UnsafeAtomics based kernels, operating on raw pointers
+
+@kernel function unsafe_atomics_add!(hist)
+    i = @index(Global, Linear)
+    j = (i - 1) % length(hist) + 1
+    UnsafeAtomics.add!(pointer(hist, j), one(eltype(hist)))
+end
+
+@kernel function unsafe_atomics_minmax!(A)
+    i = @index(Global, Linear)
+    T = eltype(A)
+    UnsafeAtomics.max!(pointer(A, 1), T(i))
+    UnsafeAtomics.min!(pointer(A, 2), T(i))
+end
+
+@kernel function unsafe_atomics_ops!(A, B)
+    i = @index(Global, Linear)
+    T = eltype(A)
+    p = pointer(A, i)
+    UnsafeAtomics.store!(p, T(i))
+    old, new = UnsafeAtomics.modify!(p, +, T(1))
+    (; success) = UnsafeAtomics.cas!(p, new, T(2) * new)
+    if success
+        old = UnsafeAtomics.xchg!(p, old)
+    end
+    B[i] = UnsafeAtomics.load(p)
+end
+
+@kernel function unsafe_atomics_add_ordered!(hist, ordering)
+    i = @index(Global, Linear)
+    j = (i - 1) % length(hist) + 1
+    UnsafeAtomics.add!(pointer(hist, j), one(eltype(hist)), ordering)
+end
+
+@kernel function unsafe_atomics_load_store_ordered!(A, B)
+    i = @index(Global, Linear)
+    v = UnsafeAtomics.load(pointer(B, i), UnsafeAtomics.acquire)
+    UnsafeAtomics.store!(pointer(A, i), v, UnsafeAtomics.release)
+end
+
+# Non-blocking message passing: workitem 1 publishes data guarded by a flag
+# with release/acquire fences; observers must see the data if they see the flag.
+@kernel function unsafe_atomics_fence!(data, flag, observed)
+    i = @index(Global, Linear)
+    T = eltype(data)
+    @inbounds if i == 1
+        data[1] = T(42)
+        UnsafeAtomics.fence(UnsafeAtomics.release)
+        UnsafeAtomics.store!(pointer(flag, 1), one(T), UnsafeAtomics.monotonic)
+    else
+        f = UnsafeAtomics.load(pointer(flag, 1), UnsafeAtomics.monotonic)
+        UnsafeAtomics.fence(UnsafeAtomics.acquire)
+        observed[i] = f == one(T) ? data[1] : T(-1)
+    end
+end
+
+@kernel function unsafe_atomics_syncscope!(A, hist)
+    i = @index(Global, Linear)
+    T = eltype(A)
+    # contended, system scope
+    j = (i - 1) % length(hist) + 1
+    UnsafeAtomics.add!(pointer(hist, j), one(T), UnsafeAtomics.seq_cst, UnsafeAtomics.none)
+    # uncontended, singlethread scope
+    p = pointer(A, i)
+    UnsafeAtomics.store!(p, T(i), UnsafeAtomics.monotonic, UnsafeAtomics.singlethread)
+    UnsafeAtomics.fence(UnsafeAtomics.seq_cst, UnsafeAtomics.singlethread)
+    UnsafeAtomics.add!(p, one(T), UnsafeAtomics.monotonic, UnsafeAtomics.singlethread)
+end
+
+function atomics_testsuite(backend, ArrayT)
+    if !KernelAbstractions.supports_atomics(backend())
+        @test_skip "Backend does not support atomics"
+        return
+    end
+
+    # Float atomics on the CPU backend need the SPV_EXT_shader_atomic_float_{add,min_max}
+    # extensions, which the POCL compiler permits based on cl_ext_float_atomics
+    eltypes = [Int32, UInt32, Float32]
+    KernelAbstractions.supports_float64(backend()) && push!(eltypes, Float64)
+
+    @testset "Atomix" begin
+        @testset "atomic add ($T)" for T in eltypes
+            hist = ArrayT(zeros(T, 32))
+            atomix_add!(backend())(hist, ndrange = 1024)
+            synchronize(backend())
+            @test all(Array(hist) .== T(1024 ÷ 32))
+        end
+
+        @testset "atomic max/min ($T)" for T in eltypes
+            A = ArrayT(zeros(T, 1))
+            atomix_max!(backend())(A, ndrange = 1024)
+            synchronize(backend())
+            @test Array(A)[1] == T(1024)
+
+            A = ArrayT(fill(typemax(T), 1))
+            atomix_min!(backend())(A, ndrange = 1024)
+            synchronize(backend())
+            @test Array(A)[1] == T(1)
+        end
+
+        @testset "atomic load/store" begin
+            A = ArrayT(zeros(Int32, 256))
+            B = ArrayT{Int32}(collect(Int32, 1:256))
+            atomix_load_store!(backend())(A, B, ndrange = 256)
+            synchronize(backend())
+            @test Array(A) == 1:256
+        end
+
+        @testset "atomicswap" begin
+            A = ArrayT(fill(Int32(-1), 256))
+            B = ArrayT(zeros(Int32, 256))
+            atomix_swap!(backend())(A, B, ndrange = 256)
+            synchronize(backend())
+            @test Array(A) == 1:256
+            @test all(Array(B) .== -1)
+        end
+
+        @testset "atomicreplace" begin
+            A = ArrayT(zeros(Int32, 256))
+            success = ArrayT(zeros(Bool, 256))
+            atomix_replace!(backend())(A, success, ndrange = 256)
+            synchronize(backend())
+            @test Array(A) == 1:256
+            @test all(Array(success))
+        end
+
+        @testset "orderings" begin
+            A = ArrayT(zeros(Int32, 256))
+            B = ArrayT(zeros(Int32, 256))
+            atomix_ordered!(backend())(A, B, ndrange = 256)
+            synchronize(backend())
+            @test Array(A) == (1:256) .+ 3
+            @test Array(B) == (1:256) .+ 3
+        end
+    end
+
+    @testset "UnsafeAtomics" begin
+        @testset "atomic add ($T)" for T in eltypes
+            hist = ArrayT(zeros(T, 32))
+            unsafe_atomics_add!(backend())(hist, ndrange = 1024)
+            synchronize(backend())
+            @test all(Array(hist) .== T(1024 ÷ 32))
+        end
+
+        @testset "atomic max/min ($T)" for T in eltypes
+            A = ArrayT(T[0, typemax(T)])
+            unsafe_atomics_minmax!(backend())(A, ndrange = 1024)
+            synchronize(backend())
+            @test Array(A) == T[1024, 1]
+        end
+
+        @testset "store/modify/cas/xchg/load" begin
+            A = ArrayT(zeros(Int32, 256))
+            B = ArrayT(zeros(Int32, 256))
+            unsafe_atomics_ops!(backend())(A, B, ndrange = 256)
+            synchronize(backend())
+            # store i, modify + 1, cas to 2(i + 1), xchg back to i
+            @test Array(A) == 1:256
+            @test Array(B) == 1:256
+        end
+
+        @testset "ordering $ordering" for ordering in (
+                UnsafeAtomics.monotonic, UnsafeAtomics.acquire, UnsafeAtomics.release,
+                UnsafeAtomics.acq_rel, UnsafeAtomics.seq_cst,
+            )
+            hist = ArrayT(zeros(Int32, 32))
+            unsafe_atomics_add_ordered!(backend())(hist, ordering, ndrange = 1024)
+            synchronize(backend())
+            @test all(Array(hist) .== 32)
+        end
+
+        @testset "ordered load/store" begin
+            A = ArrayT(zeros(Int32, 256))
+            B = ArrayT(collect(Int32, 1:256))
+            unsafe_atomics_load_store_ordered!(backend())(A, B, ndrange = 256)
+            synchronize(backend())
+            @test Array(A) == 1:256
+        end
+
+        @testset "fences" begin
+            data = ArrayT(zeros(Int32, 1))
+            flag = ArrayT(zeros(Int32, 1))
+            observed = ArrayT(zeros(Int32, 1024))
+            unsafe_atomics_fence!(backend())(data, flag, observed, ndrange = 1024)
+            synchronize(backend())
+            # observers either did not see the flag (-1) or must see the data
+            @test all(x -> x == -1 || x == 42, Array(observed)[2:end])
+        end
+
+        @testset "syncscopes" begin
+            A = ArrayT(zeros(Int32, 1024))
+            hist = ArrayT(zeros(Int32, 32))
+            unsafe_atomics_syncscope!(backend())(A, hist, ndrange = 1024)
+            synchronize(backend())
+            @test Array(A) == (1:1024) .+ 1
+            @test all(Array(hist) .== 32)
+        end
+    end
+    return
+end