diff --git a/Project.toml b/Project.toml index c1d103334..a6149ae59 100644 --- a/Project.toml +++ b/Project.toml @@ -41,7 +41,7 @@ StaticArraysExt = "StaticArrays" [compat] Adapt = "0.4, 1.0, 2.0, 3.0, 4" -Atomix = "0.1, 1" +Atomix = "1.2.1" EnzymeCore = "0.7, 0.8.1" GPUCompiler = "2" InteractiveUtils = "1.6" diff --git a/docs/src/examples/atomix.md b/docs/src/examples/atomix.md index 5bcd3998d..453145231 100644 --- a/docs/src/examples/atomix.md +++ b/docs/src/examples/atomix.md @@ -69,3 +69,32 @@ simshow(out_fixed) This image is free of artifacts. ![Resulting image is correct.](../assets/atomix_correct.png) + +## Supported operations + +`@atomic` is lowered to the atomic intrinsics of the backend in use, so which +operations and element types work depends on the backend. The following are +supported on every backend that reports `KernelAbstractions.supports_atomics(backend) == true`: + +| Operation | `Int32`, `UInt32`, `Int64`, `UInt64` | `Float32`, `Float64`[^1] | +|:-----------------------------------------------|:------------------------------------:|:------------------------:| +| `@atomic A[i]`, `@atomic A[i] = x` | ✓ | ✓ | +| `@atomicswap A[i] = x` | ✓ | ✓ | +| `@atomic A[i] += x`, `@atomic A[i] -= x` | ✓ | ✓ | +| `@atomic A[i] &= x`, `@atomic A[i] \|= x`, `@atomic A[i] ⊻= x` | ✓ | | +| `@atomic max(A[i], x)`, `@atomic min(A[i], x)` | ✓ | ✓ | +| `@atomicreplace A[i] expected => desired` | ✓ | ✓ | + +[^1]: `Float64` additionally requires `KernelAbstractions.supports_float64(backend) == true`. + +Not every backend has a native instruction for every entry in this table; for +example, CUDA has no floating-point atomic `max`/`min`. Atomix 1.2.1 and later, +which KernelAbstractions requires, fill those gaps with a compare-and-swap loop, so +the operations above work everywhere, but expect the emulated ones to be slower +under contention. Other update functions, `@atomic f(A[i], x)` for an arbitrary +binary `f`, take the same compare-and-swap path. + +A memory ordering such as `@atomic :monotonic A[i] += x` is accepted everywhere. +On the CPU backend it is honored throughout; on GPU backends the read-modify-write +operations run with the backend's default ordering, and only atomic loads and +stores follow the requested one. diff --git a/src/pocl/compiler/compilation.jl b/src/pocl/compiler/compilation.jl index 1fbc06dfe..b88fa56c9 100644 --- a/src/pocl/compiler/compilation.jl +++ b/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,7 +198,11 @@ 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 @@ -162,8 +210,12 @@ end 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 diff --git a/src/pocl/nanoOpenCL.jl b/src/pocl/nanoOpenCL.jl index eec955990..2d196c9fe 100644 --- a/src/pocl/nanoOpenCL.jl +++ b/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) diff --git a/test/atomics.jl b/test/atomics.jl new file mode 100644 index 000000000..32e517c45 --- /dev/null +++ b/test/atomics.jl @@ -0,0 +1,49 @@ +using KernelAbstractions +using KernelAbstractions: @atomic +using Test + +@kernel function atomic_add_kernel!(hist) + i = @index(Global, Linear) + j = (i - 1) % length(hist) + 1 + @inbounds @atomic hist[j] += one(eltype(hist)) +end + +@kernel function atomic_max_kernel!(A) + i = @index(Global, Linear) + @inbounds @atomic max(A[1], eltype(A)(i)) +end + +@kernel function atomic_min_kernel!(A) + i = @index(Global, Linear) + @inbounds @atomic min(A[1], eltype(A)(i)) +end + +function atomics_testsuite(backend, ArrayT) + if !KernelAbstractions.supports_atomics(backend()) + @test_skip "Backend does not support atomics" + return + end + + eltypes = [Int32, UInt32, Float32] + KernelAbstractions.supports_float64(backend()) && push!(eltypes, Float64) + + @testset "atomic add ($T)" for T in eltypes + hist = ArrayT(zeros(T, 32)) + atomic_add_kernel!(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)) + atomic_max_kernel!(backend())(A; ndrange = 1024) + synchronize(backend()) + @test Array(A)[1] == T(1024) + + A = ArrayT(fill(typemax(T), 1)) + atomic_min_kernel!(backend())(A; ndrange = 1024) + synchronize(backend()) + @test Array(A)[1] == T(1) + end + return +end diff --git a/test/runtests.jl b/test/runtests.jl index 221d35bfc..c6bacc7cd 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -13,6 +13,19 @@ KernelAbstractions.versioninfo(POCLBackend()) import KernelAbstractions.POCL: POCL, @opencl, @device_code_llvm +@testset "POCL float atomics" begin + # pocl's CPU device natively supports float add and min/max atomics in both global + # and local memory, so the SPIR-V extensions guarding them must be permitted + dev = POCL.device() + exts = split(POCL.default_spirv_extensions(dev), ",") + @test "+SPV_EXT_shader_atomic_float_add" in exts + @test "+SPV_EXT_shader_atomic_float_min_max" in exts + @test dev.half_fp_atomic_capabilities == 0 + # an explicit list overrides the device-derived default + config = POCL.compiler_config(dev; extensions = "+SPV_KHR_expect_assume") + @test config.target.extensions == "+SPV_KHR_expect_assume" +end + @testset "POCL compilation cache" begin mod = @eval module $(gensym()) @noinline child() = return diff --git a/test/testsuite.jl b/test/testsuite.jl index 79ee615ae..c647b14e6 100644 --- a/test/testsuite.jl +++ b/test/testsuite.jl @@ -27,6 +27,7 @@ end include("test.jl") +include("atomics.jl") include("hostinterface.jl") include("localmem.jl") include("private.jl") @@ -46,6 +47,10 @@ function testsuite(backend, backend_str, backend_mod, AT, DAT; skip_tests = Set{ unittest_testsuite(backend, backend_str, backend_mod, DAT; skip_tests) end + @conditional_testset "Atomics" skip_tests begin + atomics_testsuite(backend, AT) + end + @conditional_testset "SpecialFunctions" skip_tests begin specialfunctions_testsuite(backend) end