Fill GPUCompiler's SPIR-V atomics descriptor from the device

Project.toml

@@ -39,14 +39,14 @@ Adapt = "4"
 CEnum = "0.4, 0.5"
 ExprTools = "0.1"
 GPUArrays = "11.5.14"
-GPUCompiler = "2.12"
+GPUCompiler = "2.13"
 GPUToolbox = "3.3.1"
 KernelAbstractions = "0.9.39"
 LLVM = "10"
 NEO_jll = "=26.18.38308"
 PrecompileTools = "1"
 Preferences = "1"
-SPIRVIntrinsics = "1"
+SPIRVIntrinsics = "1.4"
 SPIRV_LLVM_Backend_jll = "23"
 SPIRV_LLVM_Translator_jll = "23"
 SPIRV_Tools_jll = "2025.4.0"

docs/src/api/kernels.md

@@ -234,6 +234,17 @@ oneAPI.atomic_min!(ptr, value)
 oneAPI.atomic_max!(ptr, value)
 ```
 
+Atomic operations on 8- and 16-bit values (e.g., `Int8` or `Float16`, through
+UnsafeAtomics.jl, Atomix.jl or KernelAbstractions.jl) are generally implemented by the
+compiler as 32-bit compare-and-swap loops on the aligned 4-byte word that contains the value
+(a half-precision addition can use a native instruction instead, if the device supports it).
+Arrays allocated by oneAPI.jl are padded to a multiple of 4 bytes so that this word is always
+part of the allocation. When using such atomics on memory that oneAPI.jl did not allocate,
+e.g. a pointer obtained from another library, the whole aligned 4-byte word containing every
+element you access atomically must be accessible to the device. Other elements in that word,
+such as adjacent elements of the same array, must not be modified with plain stores while an
+atomic operation on the word may be in flight; access them atomically too.
+
 Example histogram kernel:
 
 ```julia

lib/level-zero/device.jl

@@ -1,4 +1,4 @@
-export ZeDevice, properties, compute_properties, module_properties, memory_properties, memory_access_properties, cache_properties, image_properties, p2p_properties
+export ZeDevice, properties, compute_properties, module_properties, float_atomic_properties, memory_properties, memory_access_properties, cache_properties, image_properties, p2p_properties
 
 struct ZeDevice
     handle::ze_device_handle_t
@@ -105,6 +105,26 @@ function module_properties(dev::ZeDevice)
     )
 end
 
+# the floating-point atomic operations the device supports (`ze_device_fp_atomic_ext_flags_t`
+# per precision), or `nothing` if the driver doesn't implement ZE_extension_float_atomics
+function float_atomic_properties(dev::ZeDevice)
+    haskey(extension_properties(dev.driver), ZE_FLOAT_ATOMICS_EXT_NAME) || return nothing
+
+    props_ref = Ref(ze_device_module_properties_t())
+    atomic_props_ref = Ref(ze_float_atomic_ext_properties_t())
+    GC.@preserve props_ref atomic_props_ref begin
+        link_extensions(props_ref, atomic_props_ref)
+        zeDeviceGetModuleProperties(dev, props_ref)
+    end
+
+    props = atomic_props_ref[]
+    return (
+        fp16flags=props.fp16Flags,
+        fp32flags=props.fp32Flags,
+        fp64flags=props.fp64Flags,
+    )
+end
+
 function memory_properties(dev::ZeDevice)
     count_ref = Ref{UInt32}(0)
     zeDeviceGetMemoryProperties(dev, count_ref, C_NULL)

src/compiler/compilation.jl

@@ -264,6 +264,32 @@ function _driver_supports_bfloat16_spirv(dev=device())
     end
 end
 
+# The atomic operations that GPUCompiler may select SPIR-V instructions for; the others are
+# implemented with compare-and-swap loops.
+function _atomics_capabilities(dev, properties; supports_fp16, supports_fp64)
+    int64 = properties.flags & oneL0.ZE_DEVICE_MODULE_FLAG_INT64_ATOMICS ==
+            oneL0.ZE_DEVICE_MODULE_FLAG_INT64_ATOMICS
+
+    fp_atomics = oneL0.float_atomic_properties(dev)
+    fp_atomics === nothing && return SPIRVAtomics(; int64)
+    has_add(flags, flag) = flags & flag == flag
+    global_add = oneL0.ZE_DEVICE_FP_ATOMIC_EXT_FLAG_GLOBAL_ADD
+    local_add = oneL0.ZE_DEVICE_FP_ATOMIC_EXT_FLAG_LOCAL_ADD
+
+    # IGC rejects (and exits the process on) SPIR-V modules with half-precision atomic
+    # additions on devices that don't report them, e.g. Xe-LP. Xe-LP does report
+    # double-precision atomics without supporting double-precision arithmetic, though.
+    return SPIRVAtomics(;
+        int64,
+        fadd_f16_global = supports_fp16 && has_add(fp_atomics.fp16flags, global_add),
+        fadd_f16_local  = supports_fp16 && has_add(fp_atomics.fp16flags, local_add),
+        fadd_f32_global = has_add(fp_atomics.fp32flags, global_add),
+        fadd_f32_local  = has_add(fp_atomics.fp32flags, local_add),
+        fadd_f64_global = supports_fp64 && has_add(fp_atomics.fp64flags, global_add),
+        fadd_f64_local  = supports_fp64 && has_add(fp_atomics.fp64flags, local_add),
+    )
+end
+
 @noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false, kwargs...)
     properties = oneL0.module_properties(dev)
     supports_fp16 = properties.fp16flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 == oneL0.ZE_DEVICE_MODULE_FLAG_FP16
@@ -273,13 +299,11 @@ end
     # Khronos translator; the rolling stack uses the LLVM SPIR-V back-end. GPUCompiler picks
     # the tool from the target's `backend` field and loads the JLL lazily, so both can be
     # listed as deps and the choice is made here at compile time. Either way the extensions
-    # below have to be declared explicitly: without SPV_EXT_shader_atomic_float_add,
-    # floating-point atomic operations fail to translate ("The atomic float instruction
-    # requires ... SPV_EXT_shader_atomic_float_add").
+    # below have to be declared explicitly; the ones that atomics need are added by
+    # GPUCompiler based on `atomics`.
     # TODO: emit printf format strings in constant memory
     extensions = String[
         "SPV_EXT_relaxed_printf_string_address_space",
-        "SPV_EXT_shader_atomic_float_add",
     ]
     if oneL0.LTS[]
         backend = :khronos
@@ -299,8 +323,11 @@ end
     end
     extensions_str = join(map(ext -> "+$ext", extensions), ",")
 
+    atomics = _atomics_capabilities(dev, properties; supports_fp16, supports_fp64)
+
     # create GPUCompiler objects
-    target = SPIRVCompilerTarget(; backend, extensions = extensions_str, supports_fp16, supports_fp64, supports_bfloat16,
+    target = SPIRVCompilerTarget(; backend, extensions = extensions_str, atomics,
+                                   supports_fp16, supports_fp64, supports_bfloat16,
                                    driver = :intel, kwargs...)
     params = oneAPICompilerParams()
     CompilerConfig(target, params; kernel, name, always_inline)

src/compiler/execution.jl

@@ -4,7 +4,7 @@ export @oneapi, zefunction, kernel_convert
 ## high-level @oneapi interface
 
 const MACRO_KWARGS = [:launch]
-const COMPILER_KWARGS = [:kernel, :name, :always_inline]
+const COMPILER_KWARGS = [:kernel, :name, :always_inline, :atomics]
 const LAUNCH_KWARGS = [:groups, :items, :queue]
 
 """
@@ -25,6 +25,14 @@ launches the kernel on the GPU.
 - `kernel::Bool=false`: Whether to compile as a kernel (true) or device function (false)
 - `name::Union{String,Nothing}=nothing`: Explicit name for the kernel
 - `always_inline::Bool=false`: Whether to always inline device functions
+- `atomics::SPIRVAtomics`: Override the atomic operations GPUCompiler may select SPIR-V
+  instructions for, e.g. `atomics=oneAPI.SPIRVAtomics(fadd_f32_global=true)`. This replaces
+  the whole descriptor that is otherwise derived from the device and toolchain. Enabling
+  operations the device or toolchain doesn't support can make compilation fail, or terminate
+  the process from within the driver's compiler (as IGC does for half-precision atomic
+  additions on Xe-LP). Disabled floating-point operations are implemented with integer
+  compare-and-swap loops instead, and disabling `int64` makes 64-bit atomic operations an
+  error.
 
 ## Launch Keywords (runtime)
 - `groups`: Number of workgroups (required). Can be an integer or tuple.

src/pool.jl

@@ -56,8 +56,18 @@ function _maybe_gc(dev, bytes)
     end
 end
 
+# GPUCompiler implements 8- and 16-bit atomics with 32-bit atomics on the containing aligned
+# word, so every buffer has to extend to a multiple of 4 bytes and be aligned accordingly.
+# The array's dimensions stay as requested; only the buffer gets bigger.
+function atomic_allocation(bytes::Int, alignment::Int)
+    padded = bytes + (-bytes & 3)
+    padded < bytes && throw(ArgumentError("allocation of $bytes bytes is too large"))
+    return padded, max(alignment, 4)
+end
+
 function allocate(::Type{oneL0.DeviceBuffer}, ctx, dev, bytes::Int, alignment::Int)
     bytes == 0 && return oneL0.DeviceBuffer(ZE_NULL, bytes, ctx, dev)
+    bytes, alignment = atomic_allocation(bytes, alignment)
 
     _maybe_gc(dev, bytes)
     buf = device_alloc(ctx, dev, bytes, alignment)
@@ -69,6 +79,7 @@ end
 
 function allocate(::Type{oneL0.SharedBuffer}, ctx, dev, bytes::Int, alignment::Int)
     bytes == 0 && return oneL0.SharedBuffer(ZE_NULL, bytes, ctx, dev)
+    bytes, alignment = atomic_allocation(bytes, alignment)
 
     # TODO: support cross-device shared buffers (by setting `dev=nothing`)
 
@@ -82,6 +93,7 @@ end
 
 function allocate(::Type{oneL0.HostBuffer}, ctx, dev, bytes::Int, alignment::Int)
     bytes == 0 && return oneL0.HostBuffer(ZE_NULL, bytes, ctx)
+    bytes, alignment = atomic_allocation(bytes, alignment)
     buf = host_alloc(ctx, bytes, alignment)
     # Host USM must be made resident on the device, exactly like the device and shared
     # allocations above. A GPU kernel that reads a non-resident host buffer can take a

test/Project.toml

@@ -21,6 +21,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 libigc_jll = "94295238-5935-5bd7-bb0f-b00942e9bdd5"
 oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 oneAPI_Support_jll = "b049733a-a71d-5ed3-8eba-7d323ac00b36"

test/device/intrinsics.jl

@@ -1,3 +1,5 @@
+using UnsafeAtomics
+
 @testset "work items" begin
     @on_device get_work_dim() |> sink
 
@@ -277,9 +279,6 @@ end
 # @testset "atomics (low level)" begin
 
     @testset "atomic_add($T)" for T in [Int32, UInt32, Float32]
-        if oneAPI.is_integrated() && T == Float32
-            continue
-        end
         a = oneArray([zero(T)])
 
         function kernel(a, b)
@@ -292,9 +291,6 @@ end
     end
 
     @testset "atomic_sub($T)" for T in [Int32, UInt32, Float32]
-        if oneAPI.is_integrated() && T == Float32
-            continue
-        end
         a = oneArray([T(256)])
 
         function kernel(a, b)
@@ -331,9 +327,6 @@ end
     end
 
     @testset "atomic_min($T)" for T in [Int32, UInt32, Float32]
-        if oneAPI.is_integrated() && T == Float32
-            continue
-        end
         a = oneArray([T(256)])
 
         function kernel(a, T)
@@ -347,9 +340,6 @@ end
     end
 
     @testset "atomic_max($T)" for T in [Int32, UInt32, Float32]
-        if oneAPI.is_integrated() && T == Float32
-            continue
-        end
         a = oneArray([zero(T)])
 
         function kernel(a, T)
@@ -417,9 +407,6 @@ end
     end
 
     @testset "atomic_xchg($T)" for T in [Int32, UInt32, Float32]
-        if oneAPI.is_integrated() && T == Float32
-            continue
-        end
         a = oneArray([zero(T)])
 
         function kernel(a, b)
@@ -433,6 +420,55 @@ end
 
 # end
 
+@testset "atomics keyword" begin
+    function kernel(a, b)
+        oneAPI.atomic_add!(pointer(a), b)
+        return
+    end
+    a = oneArray(Float32[0])
+    tt = Tuple{typeof(oneAPI.kernel_convert(a)), Float32}
+
+    spirv = sprint(io -> oneAPI.code_spirv(io, kernel, tt; kernel=true))
+    fp_atomics = oneL0.float_atomic_properties(device())
+    if fp_atomics !== nothing &&
+       fp_atomics.fp32flags & oneL0.ZE_DEVICE_FP_ATOMIC_EXT_FLAG_GLOBAL_ADD != 0
+        @test occursin("OpAtomicFAddEXT", spirv)
+    end
+
+    # without the capability, a compare-and-swap loop is used
+    atomics = oneAPI.SPIRVAtomics()
+    spirv = sprint(io -> oneAPI.code_spirv(io, kernel, tt; kernel=true, atomics))
+    @test !occursin("OpAtomicFAddEXT", spirv)
+    @test occursin("OpAtomicCompareExchange", spirv)
+
+    @oneapi items=256 atomics=atomics kernel(a, 1f0)
+    @test Array(a)[1] == 256
+end
+
+@testset "8- and 16-bit atomics" begin
+    # these are implemented on the containing 32-bit word, so neighbouring elements contend,
+    # and the last element of an odd-sized array needs the allocation to be padded
+    function kernel(a, x)
+        i = (get_global_id() - 1) % length(a) + 1
+        UnsafeAtomics.modify!(pointer(a, i), +, x, UnsafeAtomics.monotonic,
+                              UnsafeAtomics.device)
+        return
+    end
+
+    types = [(Int8, 5), (UInt16, 3)]
+    float16_supported && push!(types, (Float16, 3))
+    @testset "$T" for (T, n) in types
+        a = oneArray(zeros(T, n))
+        @test length(a) == n
+        @test sizeof(a.data[]) % 4 == 0
+
+        # 1020 work-items, so that every element is incremented as often
+        @oneapi items=204 groups=5 kernel(a, one(T))
+        k = 1020 ÷ n
+        @test Array(a) == fill(T <: Integer ? k % T : T(k), n)
+    end
+end
+
 
 
 ############################################################################################

test/level-zero.jl

@@ -47,6 +47,7 @@ show(devnull, MIME("text/plain"), dev)
 properties(dev)
 compute_properties(dev)
 module_properties(dev)
+float_atomic_properties(dev)
 memory_properties(dev)
 memory_access_properties(dev)
 cache_properties(dev)