Add atomic fences for Metal 3.2+

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"

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)),

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

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