Demote unordered atomic accesses for SPIR-V targets

Project.toml

@@ -1,6 +1,6 @@
 name = "GPUCompiler"
 uuid = "61eb1bfa-7361-4325-ad38-22787b887f55"
-version = "2.5.3"
+version = "2.5.4"
 authors = ["Tim Besard <tim.besard@gmail.com>"]
 
 [workspace]

src/irgen.jl

@@ -354,6 +354,29 @@ function inline_unreachable_control_flow!(@nospecialize(job::CompilerJob), mod::
     return changed
 end
 
+# demote unordered atomic loads and stores to plain ones
+#
+# Julia marks accesses to heap references `unordered` so that a read racing with the GC, or
+# with another thread's write, cannot observe a torn pointer. There is no device GC and no
+# such race for GPUCompiler to protect against, so the ordering carries no meaning here, but
+# not every back-end can express it: SPIR-V's OpAtomicLoad/OpAtomicStore only take scalar
+# integer or floating-point operands, so the Khronos translator turns an `unordered` load of a
+# pointer into an invalid pointer-typed atomic that consumers reject (Intel's compiler fails
+# with an undefined `__spirv_AtomicLoad(long**, int, int)`), and AIR has no atomic load or
+# store instructions at all. Run after optimization, where dropping the ordering cannot
+# enable new transformations; stronger orderings are left intact.
+function demote_unordered_atomics!(mod::LLVM.Module)
+    changed = false
+    for f in functions(mod), bb in blocks(f), inst in instructions(bb)
+        (inst isa LLVM.LoadInst || inst isa LLVM.StoreInst) || continue
+        is_atomic(inst) || continue
+        ordering(inst) == LLVM.API.LLVMAtomicOrderingUnordered || continue
+        ordering!(inst, LLVM.API.LLVMAtomicOrderingNotAtomic)
+        changed = true
+    end
+    return changed
+end
+
 # lower `trap` to a clean return to get rid of `unreachable` and `noreturn`
 #
 # this is for compatibility with back-ends that don't support (SPIR-V) or have

src/metal.jl

@@ -645,23 +645,6 @@ function finish_ir!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
     return functions(mod)[entry_fn]
 end
 
-# Julia marks heap-reference accesses `unordered` so raced pointer reads cannot produce
-# values out of thin air. GPUCompiler has no device GC, and Metal's supported atomics
-# are already `air.atomic.*` intrinsics, so raw unordered LLVM loads/stores are unnecessary.
-# Demote them after optimization, where the weaker semantics cannot enable new transforms;
-# leave stronger orderings intact.
-function demote_unordered_atomics!(mod::LLVM.Module)
-    changed = false
-    for f in functions(mod), bb in blocks(f), inst in instructions(bb)
-        (inst isa LLVM.LoadInst || inst isa LLVM.StoreInst) || continue
-        is_atomic(inst) || continue
-        ordering(inst) == LLVM.API.LLVMAtomicOrderingUnordered || continue
-        ordering!(inst, LLVM.API.LLVMAtomicOrderingNotAtomic)
-        changed = true
-    end
-    return changed
-end
-
 # lowering of LLVM IR to AIR-compatible IR
 #
 # Metal does not have an LLVM back-end, so the lowering of target-independent LLVM IR into
@@ -676,7 +659,7 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
     # Metal.malloc uses.
     rewrite_generic_null_selects!(mod)
 
-    # AIR does not support LLVM atomic load/store instructions
+    # AIR does not support LLVM atomic load/store instructions (see `demote_unordered_atomics!`)
     demote_unordered_atomics!(mod)
 
     # strip device-side `trap`s and rewrite `unreachable` into clean returns (#433, #370). this

src/spirv.jl

@@ -109,6 +109,10 @@ function finish_ir!(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module,
     # OpUnreachable (UB if reached), which PoCL and friends handle poorly.
     lower_unreachable_control_flow!(job, mod)
 
+    # SPIR-V cannot express atomic loads and stores of pointers, which is what Julia's
+    # `unordered` heap-reference accesses become; the orderings serve no purpose on device
+    demote_unordered_atomics!(mod)
+
     # convert the kernel state argument to a byval reference
     if job.config.kernel
         state = kernel_state_type(job)

test/spirv.jl

@@ -214,6 +214,38 @@ end
     end
 end
 
+@testset "unordered atomic demotion" begin
+    # Julia's `unordered` heap-reference accesses cannot be expressed in SPIR-V when they
+    # involve pointers (OpAtomicLoad/OpAtomicStore take scalars only): the translator would
+    # emit an invalid pointer-typed atomic. They carry no meaning without a device GC, so
+    # `demote_unordered_atomics!` turns them into plain accesses.
+    mod = @eval module $(gensym())
+        function kernel(p::Ptr{Ptr{Int}}, q::Ptr{Ptr{Int}})
+            x = Core.Intrinsics.atomic_pointerref(p, :unordered)
+            Core.Intrinsics.atomic_pointerset(q, x, :unordered)
+            return
+        end
+    end
+    tt = Tuple{Ptr{Ptr{Int}}, Ptr{Ptr{Int}}}
+
+    @test @filecheck begin
+        @check_label "define spir_kernel void @_Z6kernel"
+        @check_not "load atomic"
+        @check_not "store atomic"
+        @check "ret void"
+        SPIRV.code_llvm(mod.kernel, tt; backend, kernel=true)
+    end
+
+    # the SPIR-V is validated by the helper, so an invalid pointer-typed atomic would fail here
+    @test @filecheck begin
+        @check "OpEntryPoint Kernel %[[KERNEL:[^ ]+]]"
+        @check "%[[KERNEL]] = OpFunction %void None"
+        @check_not "OpAtomicLoad"
+        @check_not "OpAtomicStore"
+        SPIRV.code_native(mod.kernel, tt; backend, kernel=true)
+    end
+end
+
 @testset "inlining of throwing callees" begin
     mod = @eval module $(gensym())
         @noinline function guard(x)