diff --git a/Project.toml b/Project.toml index ae9646f1..f35c5ee3 100644 --- a/Project.toml +++ b/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" diff --git a/lib/intrinsics/Project.toml b/lib/intrinsics/Project.toml index d552c7a2..6b9079fc 100644 --- a/lib/intrinsics/Project.toml +++ b/lib/intrinsics/Project.toml @@ -1,7 +1,7 @@ name = "SPIRVIntrinsics" uuid = "71d1d633-e7e8-4a92-83a1-de8814b09ba8" authors = ["Tim Besard "] -version = "1.0.1" +version = "1.1.0" [deps] ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04" diff --git a/lib/intrinsics/src/atomic.jl b/lib/intrinsics/src/atomic.jl index 27a824bf..fbf47d09 100644 --- a/lib/intrinsics/src/atomic.jl +++ b/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) diff --git a/src/OpenCL.jl b/src/OpenCL.jl index 96e4889f..ba93d318 100644 --- a/src/OpenCL.jl +++ b/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") diff --git a/src/compiler/capabilities.jl b/src/compiler/capabilities.jl index 6d419cf1..52562bad 100644 --- a/src/compiler/capabilities.jl +++ b/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 diff --git a/src/compiler/compilation.jl b/src/compiler/compilation.jl index 3ca31795..1fb9584d 100644 --- a/src/compiler/compilation.jl +++ b/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 diff --git a/src/device/atomics.jl b/src/device/atomics.jl new file mode 100644 index 00000000..56de5f8d --- /dev/null +++ b/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 diff --git a/test/atomics.jl b/test/atomics.jl index 954ea09c..00d216e3 100644 --- a/test/atomics.jl +++ b/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