diff --git a/Project.toml b/Project.toml index 72468b30..390ed55c 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/api/kernels.md b/docs/src/api/kernels.md index 0df541a4..e82282eb 100644 --- a/docs/src/api/kernels.md +++ b/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 diff --git a/lib/level-zero/device.jl b/lib/level-zero/device.jl index 58d5b2da..8b1e9f78 100644 --- a/lib/level-zero/device.jl +++ b/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) diff --git a/src/compiler/compilation.jl b/src/compiler/compilation.jl index 34706809..33b5537b 100644 --- a/src/compiler/compilation.jl +++ b/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) diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index f129919b..1500fcb5 100644 --- a/src/compiler/execution.jl +++ b/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. diff --git a/src/pool.jl b/src/pool.jl index 8d390234..a159ed8c 100644 --- a/src/pool.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index 190beeab..beab523b 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/device/intrinsics.jl b/test/device/intrinsics.jl index bd2fadd9..6130a879 100644 --- a/test/device/intrinsics.jl +++ b/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 + ############################################################################################ diff --git a/test/level-zero.jl b/test/level-zero.jl index bd070d8f..c210021f 100644 --- a/test/level-zero.jl +++ b/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)