Add floating-point atomic add/sub and min/max intrinsics

Project.toml

@@ -42,7 +42,7 @@ Random = "1"
 Random123 = "1.7.1"
 RandomNumbers = "1.6.0"
 Reexport = "1"
-SPIRVIntrinsics = "1"
+SPIRVIntrinsics = "1.1"
 SPIRV_LLVM_Backend_jll = "22"
 SPIRV_Tools_jll = "2025.1"
 StaticArrays = "1"

lib/intrinsics/Project.toml

@@ -1,7 +1,7 @@
 name = "SPIRVIntrinsics"
 uuid = "71d1d633-e7e8-4a92-83a1-de8814b09ba8"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
-version = "1.0.1"
+version = "1.1.0"
 
 [deps]
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"

lib/intrinsics/src/atomic.jl

@@ -84,6 +84,80 @@ end
 end
 
 
+# floating-point atomics.
+#
+# Native floating-point atomic add and min/max come from the SPV_EXT_shader_atomic_float_add
+# and SPV_EXT_shader_atomic_float_min_max extensions, emitted as SPIR-V wrapper builtins like
+# the integer atomics above (OpenCL-style builtins with float arguments are not recognized by
+# the Khronos translator, and only atomic_add by the LLVM SPIR-V backend). Extension support
+# is an optional device capability (cl_ext_float_atomics), so the generic functions default
+# to a compare-and-swap loop (correct on any device with integer cmpxchg); back-ends that can
+# query the device override them to select the native version at compile time (see OpenCL.jl's
+# `has_feature`).
+for gentype in atomic_float_types, as in atomic_memory_types
+    bits = gentype == Float32 ? UInt32 : UInt64
+@eval begin
+
+@device_function atomic_add_native!(p::LLVMPtr{$gentype,$as}, val::$gentype) =
+    @builtin_ccall("__spirv_AtomicFAddEXT", $gentype,
+                   (LLVMPtr{$gentype,$as}, UInt32, UInt32, $gentype),
+                   p, UInt32(atomic_scope),
+                   UInt32(atomic_memory_semantics(Val($as))), val)
+
+# SPIR-V has no atomic float subtraction; add the negated value
+@device_function atomic_sub_native!(p::LLVMPtr{$gentype,$as}, val::$gentype) =
+    atomic_add_native!(p, -val)
+
+@device_function atomic_min_native!(p::LLVMPtr{$gentype,$as}, val::$gentype) =
+    @builtin_ccall("__spirv_AtomicFMinEXT", $gentype,
+                   (LLVMPtr{$gentype,$as}, UInt32, UInt32, $gentype),
+                   p, UInt32(atomic_scope),
+                   UInt32(atomic_memory_semantics(Val($as))), val)
+
+@device_function atomic_max_native!(p::LLVMPtr{$gentype,$as}, val::$gentype) =
+    @builtin_ccall("__spirv_AtomicFMaxEXT", $gentype,
+                   (LLVMPtr{$gentype,$as}, UInt32, UInt32, $gentype),
+                   p, UInt32(atomic_scope),
+                   UInt32(atomic_memory_semantics(Val($as))), val)
+
+end
+
+# the loops compare raw bits, not values: `==` on floats would spin forever on a stored NaN,
+# and would treat a failed exchange as successful when -0.0 compares equal to 0.0.
+for (op, expr) in [:add => :(old + val), :min => :(min(old, val)), :max => :(max(old, val))]
+    fallback = Symbol("atomic_$(op)_fallback!")
+    fn = Symbol("atomic_$(op)!")
+@eval begin
+
+@device_function @inline function $fallback(p::LLVMPtr{$gentype,$as}, val::$gentype)
+    ip = reinterpret(LLVMPtr{$bits,$as}, p)
+    cmp = Base.unsafe_load(ip, 1)
+    while true
+        old = reinterpret($gentype, cmp)
+        new = reinterpret($bits, $expr)
+        seen = atomic_cmpxchg!(ip, cmp, new)
+        seen == cmp && return old
+        cmp = seen
+    end
+end
+
+@device_function $fn(p::LLVMPtr{$gentype,$as}, val::$gentype) = $fallback(p, val)
+
+end
+end
+
+@eval begin
+
+@device_function atomic_sub_fallback!(p::LLVMPtr{$gentype,$as}, val::$gentype) =
+    atomic_add_fallback!(p, -val)
+
+@device_function atomic_sub!(p::LLVMPtr{$gentype,$as}, val::$gentype) =
+    atomic_sub_fallback!(p, val)
+
+end
+end
+
+
 # specifically typed
 
 for as in atomic_memory_types
@@ -271,15 +345,19 @@ end
 # native atomics
 # TODO: support inc/dec
 # TODO: this depends on backend support for the corresponding SPIR-V atomic
-#       operation. Floating-point arithmetic should hit the cmpxchg fallback
-#       unless a caller explicitly uses a floating-point atomic extension.
+#       operation. 64-bit integer atomics require the Int64Atomics capability.
 for (op,impl) in [(+)      => atomic_add!,
                   (-)      => atomic_sub!,
-                  (&)      => atomic_and!,
-                  (|)      => atomic_or!,
-                  (⊻)      => atomic_xor!,
                   Base.max => atomic_max!,
                   Base.min => atomic_min!]
+    @eval @inline atomic_arrayset(A::AbstractArray{T}, I::Integer, ::typeof($op),
+                                  val::T) where {T <: Union{atomic_integer_types...,
+                                                            atomic_float_types...}} =
+        $impl(pointer(A, I), val)
+end
+for (op,impl) in [(&)      => atomic_and!,
+                  (|)      => atomic_or!,
+                  (⊻)      => atomic_xor!]
     @eval @inline atomic_arrayset(A::AbstractArray{T}, I::Integer, ::typeof($op),
                                   val::T) where {T <: Union{atomic_integer_types...}} =
         $impl(pointer(A, I), val)

src/OpenCL.jl

@@ -27,6 +27,7 @@ include("device/runtime.jl")
 include("device/array.jl")
 include("device/quirks.jl")
 include("device/random.jl")
+include("device/atomics.jl")
 
 # high level implementation
 include("memory.jl")

src/compiler/capabilities.jl

@@ -40,6 +40,20 @@ end
 
 has_opencl_c_feature(dev, feat) = feat in opencl_c_features(dev)
 
+# cl_ext_float_atomics: native floating-point atomic capabilities, reported per precision as a
+# bitfield with separate global- and local-memory bits. The device intrinsics operate on either
+# address space, so require both bits.
+function has_fp_atomics(dev::cl.Device, query, caps)
+    "cl_ext_float_atomics" in dev.extensions || return false
+    try
+        supported = Ref{UInt64}(0)  # cl_device_fp_atomic_capabilities_ext
+        cl.clGetDeviceInfo(dev, query, sizeof(UInt64), supported, C_NULL)
+        return supported[] & caps == caps
+    catch
+        return false
+    end
+end
+
 # OpenCL 3.0: CL_DEVICE_OPENCL_C_ALL_VERSIONS lists every OpenCL C version the device accepts as
 # an array of cl_name_version {cl_version (4 bytes); char name[64]}. This is the query to trust:
 # the legacy CL_DEVICE_OPENCL_C_VERSION string reports "1.2" on both NVIDIA and pocl even though
@@ -90,6 +104,22 @@ const FEATURES = Feature[
     Feature(:subgroups, cl.sub_groups_supported),
     Feature(:generic_address_space,
             dev -> has_opencl_c_feature(dev, "__opencl_c_generic_address_space")),
+    Feature(:fp32_atomic_add,
+            dev -> has_fp_atomics(dev, cl.CL_DEVICE_SINGLE_FP_ATOMIC_CAPABILITIES_EXT,
+                                  cl.CL_DEVICE_GLOBAL_FP_ATOMIC_ADD_EXT |
+                                  cl.CL_DEVICE_LOCAL_FP_ATOMIC_ADD_EXT)),
+    Feature(:fp32_atomic_min_max,
+            dev -> has_fp_atomics(dev, cl.CL_DEVICE_SINGLE_FP_ATOMIC_CAPABILITIES_EXT,
+                                  cl.CL_DEVICE_GLOBAL_FP_ATOMIC_MIN_MAX_EXT |
+                                  cl.CL_DEVICE_LOCAL_FP_ATOMIC_MIN_MAX_EXT)),
+    Feature(:fp64_atomic_add,
+            dev -> has_fp_atomics(dev, cl.CL_DEVICE_DOUBLE_FP_ATOMIC_CAPABILITIES_EXT,
+                                  cl.CL_DEVICE_GLOBAL_FP_ATOMIC_ADD_EXT |
+                                  cl.CL_DEVICE_LOCAL_FP_ATOMIC_ADD_EXT)),
+    Feature(:fp64_atomic_min_max,
+            dev -> has_fp_atomics(dev, cl.CL_DEVICE_DOUBLE_FP_ATOMIC_CAPABILITIES_EXT,
+                                  cl.CL_DEVICE_GLOBAL_FP_ATOMIC_MIN_MAX_EXT |
+                                  cl.CL_DEVICE_LOCAL_FP_ATOMIC_MIN_MAX_EXT)),
 ]
 
 const FeatureSet = UInt64

src/compiler/compilation.jl

@@ -157,19 +157,42 @@ end
 const SPIRV_VERSION = v"1.4"
 
 @noinline function _compiler_config(dev, backend; kernel=true, name=nothing, always_inline=false,
-                                     sub_group_size::Union{Nothing,Int}=_sub_group_size(dev), kwargs...)
+                                     sub_group_size::Union{Nothing,Int}=_sub_group_size(dev),
+                                     extensions=nothing, kwargs...)
     supports_fp16 = "cl_khr_fp16" in dev.extensions
     supports_fp64 = "cl_khr_fp64" in dev.extensions
+    features = device_features(dev)
+
+    if backend === :opencl || !("cl_khr_il_program" in dev.extensions)
+        # programs on the OpenCL C source backend go through spirv2clc, which does not
+        # implement the SPV_EXT_shader_atomic_float_min_max capabilities; mask the features
+        # so kernels take the compare-and-swap fallback instead.
+        features &= ~(feature_bit(:fp32_atomic_min_max) | feature_bit(:fp64_atomic_min_max))
+    end
+
+    if extensions === nothing
+        # allow the SPIR-V float-atomics extensions on devices that support them, so kernels
+        # can use the native atomic instructions (see `device/atomics.jl`); OpExtension only
+        # ends up in modules that use the corresponding instructions.
+        extensions = String[]
+        if feature_supported(features, :fp32_atomic_add) ||
+           feature_supported(features, :fp64_atomic_add)
+            push!(extensions, "SPV_EXT_shader_atomic_float_add")
+        end
+        if feature_supported(features, :fp32_atomic_min_max) ||
+           feature_supported(features, :fp64_atomic_min_max)
+            push!(extensions, "SPV_EXT_shader_atomic_float_min_max")
+        end
+    end
 
     if sub_group_size !== nothing && !("cl_intel_required_subgroup_size" in dev.extensions)
         error("Device does not support cl_intel_required_subgroup_size")
     end
 
     # create GPUCompiler objects
-    target = SPIRVCompilerTarget(; version=SPIRV_VERSION, supports_fp16, supports_fp64,
-                                   validate=true, kwargs...)
-    params = OpenCLCompilerParams(; sub_group_size, features=device_features(dev),
-                                    program_backend=backend)
+    target = SPIRVCompilerTarget(; version=SPIRV_VERSION, extensions, supports_fp16,
+                                   supports_fp64, validate=true, kwargs...)
+    params = OpenCLCompilerParams(; sub_group_size, features, program_backend=backend)
     CompilerConfig(target, params; kernel, name, always_inline)
 end
 

src/device/atomics.jl

@@ -0,0 +1,31 @@
+## feature-gated floating-point atomics
+
+# SPIRVIntrinsics defaults the floating-point atomics to a compare-and-swap loop; on devices
+# that advertise native floating-point atomics (cl_ext_float_atomics), select the EXT
+# instructions instead. `has_feature` folds at compile time, so only the chosen implementation
+# is emitted, and `_compiler_config` only allows the corresponding SPIR-V extensions on
+# devices that support them.
+for (T, add_feature, min_max_feature) in
+        ((Float32, :fp32_atomic_add, :fp32_atomic_min_max),
+         (Float64, :fp64_atomic_add, :fp64_atomic_min_max)),
+    as in (AS.Workgroup, AS.CrossWorkgroup)
+@eval begin
+
+@device_override SPIRVIntrinsics.atomic_add!(p::LLVMPtr{$T,$as}, val::$T) =
+    has_feature($(QuoteNode(add_feature))) ? atomic_add_native!(p, val) :
+                                             atomic_add_fallback!(p, val)
+
+@device_override SPIRVIntrinsics.atomic_sub!(p::LLVMPtr{$T,$as}, val::$T) =
+    has_feature($(QuoteNode(add_feature))) ? atomic_sub_native!(p, val) :
+                                             atomic_sub_fallback!(p, val)
+
+@device_override SPIRVIntrinsics.atomic_min!(p::LLVMPtr{$T,$as}, val::$T) =
+    has_feature($(QuoteNode(min_max_feature))) ? atomic_min_native!(p, val) :
+                                                 atomic_min_fallback!(p, val)
+
+@device_override SPIRVIntrinsics.atomic_max!(p::LLVMPtr{$T,$as}, val::$T) =
+    has_feature($(QuoteNode(min_max_feature))) ? atomic_max_native!(p, val) :
+                                                 atomic_max_fallback!(p, val)
+
+end
+end

test/atomics.jl

@@ -48,6 +48,15 @@ function test_atomic_cas(counter::AbstractArray{T}) where T
     OpenCL.atomic_cmpxchg!(pointer(counter), zero(T), one(T))
     return
 end
+# Floating-point add/sub - use low-level API directly
+function float_add_kernel(counter::AbstractArray{T}, val::T) where T
+    OpenCL.atomic_add!(pointer(counter), val)
+    return
+end
+function float_sub_kernel(counter::AbstractArray{T}, val::T) where T
+    OpenCL.atomic_sub!(pointer(counter), val)
+    return
+end
 
 # Define atomic operations to test
 atomic_operations = [
@@ -74,13 +83,19 @@ atomic_operations = [
         continue
     end
 
+    # Float64 atomics may fall back to 64-bit cmpxchg
+    if T == Float64 && !("cl_khr_int64_base_atomics" in dev.extensions)
+        continue
+    end
+
     # Bitwise operations (only valid for integers)
     if kernel_func in [test_atomic_and, test_atomic_or, test_atomic_xor] && T <: AbstractFloat
         continue
     end
 
-    # Min/max operations (only supported for 32-bit integers in OpenCL)
-    if kernel_func in [test_atomic_min, test_atomic_max] && !(T in [Int32, UInt32])
+    # Min/max on integers is only supported for 32-bit types; floats use the native
+    # extension or the compare-and-swap fallback
+    if kernel_func in [test_atomic_min, test_atomic_max] && !(T in [Int32, UInt32] || T <: AbstractFloat)
         continue
     end
 
@@ -96,7 +111,39 @@ atomic_operations = [
 end
 
 
-@testset "atomic_add! ($T)" for T in [Float32, Float64]
+@testset "float atomics ($T)" for T in [Float32, Float64]
+    if T == Float64 && !("cl_khr_fp64" in dev.extensions)
+        continue
+    end
+    if T == Float64 && !("cl_khr_int64_base_atomics" in dev.extensions)
+        continue
+    end
+
+    a = OpenCL.zeros(T)
+    @opencl global_size=1000 float_add_kernel(a, one(T))
+    @test OpenCL.@allowscalar(a[]) == T(1000)
+
+    b = OpenCL.fill(T(1000))
+    @opencl global_size=1000 float_sub_kernel(b, one(T))
+    @test OpenCL.@allowscalar(b[]) == T(0)
+
+    # the native/fallback selection must fold at compile time, leaving only the path
+    # matching the device's capabilities
+    feature = T == Float32 ? :fp32_atomic_add : :fp64_atomic_add
+    ir = sprint() do io
+        OpenCL.code_llvm(io, float_add_kernel, Tuple{CLDeviceArray{T, 0, AS.CrossWorkgroup}, T};
+                         kernel=true, dump_module=true)
+    end
+    if OpenCL.feature_supported(dev, feature)
+        @test occursin("__spirv_AtomicFAddEXT", ir)
+        @test !occursin("__spirv_AtomicCompareExchange", ir)
+    else
+        @test occursin("__spirv_AtomicCompareExchange", ir)
+        @test !occursin("__spirv_AtomicFAddEXT", ir)
+    end
+end
+
+@testset "atomic_add builtin ($T)" for T in [Float32, Float64]
     # Float64 requires cl_khr_fp64 extension
     if T == Float64 && !("cl_khr_fp64" in cl.device().extensions)
         continue