From 997fa62292c94bf1be0326572a5e90c97ec12aa6 Mon Sep 17 00:00:00 2001 From: Christian <28689358+christiangnrd@users.noreply.github.com> Date: Thu, 16 Jul 2026 13:15:16 -0300 Subject: [PATCH] Atomic fences Co-Authored-By: Tim Besard <383068+maleadt@users.noreply.github.com> --- Project.toml | 2 +- src/device/intrinsics/atomics.jl | 4 -- src/device/intrinsics/synchronization.jl | 44 +++++++++++++++++++++- test/device/intrinsics/synchronization.jl | 45 +++++++++++++++++++++++ 4 files changed, 88 insertions(+), 7 deletions(-) diff --git a/Project.toml b/Project.toml index 06b3d4f24..2ecf2f6d3 100644 --- a/Project.toml +++ b/Project.toml @@ -50,7 +50,7 @@ CodecBzip2 = "0.8.5" Crayons = "4" ExprTools = "0.1" GPUArrays = "11.5.6" -GPUCompiler = "2" +GPUCompiler = "2.1" GPUToolbox = "3" KernelAbstractions = "0.9.38" LLVM = "7.2, 8, 9" diff --git a/src/device/intrinsics/atomics.jl b/src/device/intrinsics/atomics.jl index a82b9f0d0..d92024162 100644 --- a/src/device/intrinsics/atomics.jl +++ b/src/device/intrinsics/atomics.jl @@ -1,9 +1,5 @@ # Atomic Functions -@enum memory_order::Int32 begin - memory_order_relaxed = 0 -end - # XXX: the integers should come from some enum const atomic_memory_names = Dict( AS.Device => ("global", Int32(2)), diff --git a/src/device/intrinsics/synchronization.jl b/src/device/intrinsics/synchronization.jl index c7489a5b9..7e92224db 100644 --- a/src/device/intrinsics/synchronization.jl +++ b/src/device/intrinsics/synchronization.jl @@ -1,12 +1,28 @@ # Synchronization using CEnum -export MemoryFlags, threadgroup_barrier, simdgroup_barrier +export MemoryFlags, memory_order, thread_scope, atomic_thread_fence, + threadgroup_barrier, simdgroup_barrier + +@enum memory_order::Int32 begin + memory_order_relaxed = 0 + memory_order_acquire = 2 + memory_order_release = 3 + memory_order_acq_rel = 4 + memory_order_seq_cst = 5 +end + +@enum thread_scope::Int32 begin + thread_scope_thread = 0 + thread_scope_threadgroup = 1 + thread_scope_device = 2 + thread_scope_simdgroup = 4 +end """ MemoryFlags -Flags to set the memory synchronization behavior of threadgroup\\_barrier and simdgroup\\_barrier. +Flags to set the memory synchronization behavior of barriers and atomic fences. Possible values: @@ -39,6 +55,23 @@ end @device_function @inline simdgroup_barrier(flag=MemoryFlagNone) = ccall("extern air.simdgroup.barrier", llvmcall, Cvoid, (Cuint, Cuint, ), flag, UInt32(1)) +@device_function @inline atomic_thread_fence(flags::Union{MemoryFlags,UInt32}, + order::memory_order, + scope::thread_scope=thread_scope_device) = + atomic_thread_fence(Val(flags), Val(order), Val(scope)) + +@device_function @inline function atomic_thread_fence(flags::Val{F}, order::Val{O}, + scope::Val{S}) where {F,O,S} + @static_assert(metal_version() >= sv"3.2", + "atomic_thread_fence requires Metal 3.2 or newer.") + @static_assert(O isa memory_order, "Invalid atomic memory ordering.") + @static_assert(O === memory_order_relaxed || O === memory_order_seq_cst || metal_version() >= sv"4.1", + "Acquire, release, and acquire-release atomic_thread_fence orderings require Metal 4.1 or newer.") + @static_assert(S isa thread_scope, "Invalid atomic thread scope.") + @typed_ccall("air.atomic.fence", llvmcall, Nothing, (Int32, Int32, Int32), + flags, order, scope) +end + @doc """ threadgroup_barrier(flag=MemoryFlagNone) @@ -54,3 +87,10 @@ Synchronize all threads in a SIMD-group. Possible flags that affect the memory synchronization behavior are found in [`MemoryFlags`](@ref) """ simdgroup_barrier + +@doc """ + atomic_thread_fence(flags, order, scope=thread_scope_device) + +Order memory accesses selected by `flags` for threads in `scope`, without an execution +barrier. `flags`, `order`, and `scope` must be compile-time constants. +""" atomic_thread_fence diff --git a/test/device/intrinsics/synchronization.jl b/test/device/intrinsics/synchronization.jl index 61397ac05..b42e0ddaa 100644 --- a/test/device/intrinsics/synchronization.jl +++ b/test/device/intrinsics/synchronization.jl @@ -50,6 +50,51 @@ @test Array(buf)[1] == threads - 1 end + # TODO: Actually test for races + @testset "atomic thread fence" begin + function fence_kernel(buf, ::Val{ORDER}, ::Val{FLAGS}, ::Val{SCOPE}) where {ORDER,FLAGS,SCOPE} + Metal.atomic_thread_fence(FLAGS, ORDER, SCOPE) + buf[1] += 1 + return + end + + orders = [Metal.memory_order_relaxed, Metal.memory_order_seq_cst] + macos_version() >= v"27" && append!(orders, [Metal.memory_order_acquire, Metal.memory_order_release, Metal.memory_order_acq_rel]) + for order in orders + buf = Metal.zeros(Int32, 1) + @metal fence_kernel(buf, Val(order), + Val(Metal.MemoryFlagDevice | Metal.MemoryFlagTexture), + Val(Metal.thread_scope_simdgroup)) + @test Array(buf) == Int32[1] + end + + function fence_abi(::Core.LLVMPtr{Int32,Metal.AS.Device}) + Metal.atomic_thread_fence(Metal.MemoryFlagDevice | Metal.MemoryFlagTexture, + Metal.memory_order_seq_cst, + Metal.thread_scope_simdgroup) + return + end + ir = sprint(io -> Metal.code_llvm(io, fence_abi, + Tuple{Core.LLVMPtr{Int32,Metal.AS.Device}}; + kernel=true, metal=v"3.2", dump_module=true)) + @test occursin("@air.atomic.fence(i32, i32, i32)", ir) + @test occursin("i32 5, i32 5, i32 4", ir) + + function unavailable_fence(buf) + Metal.atomic_thread_fence(Metal.MemoryFlagDevice, Metal.memory_order_relaxed) + return + end + err = try + @metal launch=false metal=v"3.1" unavailable_fence(Metal.zeros(Int32, 1)) + nothing + catch err + err + end + @test err isa Metal.InvalidIRError + @test occursin("atomic_thread_fence requires Metal 3.2 or newer.", + sprint(showerror, err)) + end + # TODO: simdgroup barrier test end