Demote all LLVM atomic loads and stores on Metal and SPIR-V

src/irgen.jl

@@ -354,23 +354,28 @@ function inline_unreachable_control_flow!(@nospecialize(job::CompilerJob), mod::
     return changed
 end
 
-# demote unordered atomic loads and stores to plain ones
+# demote LLVM 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)
+# with another thread's write, cannot observe a torn pointer, and stores the type tag of a
+# freshly allocated object with `release` ordering so that no other thread can observe the
+# object before its header. There is no device GC and no such race for GPUCompiler to
+# protect against, so these orderings carry no meaning here, but not every back-end can
+# express them: SPIR-V's OpAtomicLoad/OpAtomicStore only take scalar integer or
+# floating-point operands, so the Khronos translator turns an atomic access 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: Apple's back-end aborts on them (`XPC_ERROR_CONNECTION_INTERRUPTED`
+# from the driver; the macOS 26 AGX compiler reports `unable to legalize instruction:
+# store release (p0)` for a type-tag store through the generic pointer the device allocator
+# returns). Run after optimization, where dropping the ordering cannot enable new
+# transformations. Device-side atomics proper go through target intrinsics, not these
+# instructions, so every remaining one is such Julia bookkeeping and gets demoted.
+function demote_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

src/metal.jl

@@ -659,8 +659,8 @@ 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 (see `demote_unordered_atomics!`)
-    demote_unordered_atomics!(mod)
+    # AIR does not support LLVM atomic load/store instructions (see `demote_atomics!`)
+    demote_atomics!(mod)
 
     # strip device-side `trap`s and rewrite `unreachable` into clean returns (#433, #370). this
     # runs post-`optimize!`, after the trap has finished serving as the optimizer guard; the pass

src/spirv.jl

@@ -110,8 +110,8 @@ function finish_ir!(job::CompilerJob{SPIRVCompilerTarget}, mod::LLVM.Module,
     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)
+    # heap-reference accesses and type-tag stores are; the orderings serve no purpose on device
+    demote_atomics!(mod)
 
     # convert the kernel state argument to a byval reference
     if job.config.kernel

test/metal.jl

@@ -1589,14 +1589,16 @@ end
     @test !occursin(r"(load|store) atomic", air)
 end
 
-@testset "unordered atomic demotion" begin
-    # Demote unordered LLVM loads/stores, but preserve stronger atomics.
+@testset "atomic demotion" begin
+    # Demote every LLVM atomic load/store, whatever its ordering or value type.
     Context() do ctx
         ir = """
         define void @f(i8** %p, i8** %q, i64* %r) {
         entry:
           %a = load atomic i8*, i8** %p unordered, align 8
           store atomic i8* %a, i8** %q unordered, align 8
+          %t = load atomic i8*, i8** %p acquire, align 8
+          store atomic i8* %t, i8** %q release, align 8
           %b = load atomic i64, i64* %r monotonic, align 8
           store atomic i64 %b, i64* %r release, align 8
           store i8* %a, i8** %p, align 8
@@ -1607,15 +1609,13 @@ end
         insts() = [i for f in functions(mod) for bb in blocks(f) for i in instructions(bb)]
         memops() = filter(i -> i isa LLVM.LoadInst || i isa LLVM.StoreInst, insts())
 
-        @test count(is_atomic, memops()) == 4
-        @test GPUCompiler.demote_unordered_atomics!(mod)
-        atomics = filter(is_atomic, memops())
-        @test length(atomics) == 2
-        @test all(i -> ordering(i) != LLVM.API.LLVMAtomicOrderingUnordered, atomics)
-        @test !occursin("unordered", string(mod))
+        @test count(is_atomic, memops()) == 6
+        @test GPUCompiler.demote_atomics!(mod)
+        @test count(is_atomic, memops()) == 0
+        @test !occursin("atomic", string(mod))
         @test (verify(mod); true)
         # idempotent
-        @test !GPUCompiler.demote_unordered_atomics!(mod)
+        @test !GPUCompiler.demote_atomics!(mod)
     end
 
     # end-to-end: Julia's own `:unordered` accesses (as codegen emits for heap-reference
@@ -1640,6 +1640,27 @@ end
     @test !occursin(r"(load|store) atomic", air)
     @test occursin(r"load i64", air)
     @test occursin(r"store i64", air)
+
+    # likewise for the `release` store of a pointer that codegen emits for an allocated
+    # object's type tag. Only Julia 1.12+ lowers `Ptr` values to LLVM pointers, which is
+    # what makes this the case Apple's back-end cannot legalize.
+    @static if VERSION >= v"1.12"
+    function tagged(p::Core.LLVMPtr{Int,1}, q::Core.LLVMPtr{Int,1})
+        x = Core.Intrinsics.atomic_pointerref(reinterpret(Ptr{Ptr{Int}}, p), :acquire)
+        Core.Intrinsics.atomic_pointerset(reinterpret(Ptr{Ptr{Int}}, q), x, :release)
+        return
+    end
+    source = methodinstance(typeof(tagged), Tuple{Core.LLVMPtr{Int,1}, Core.LLVMPtr{Int,1}},
+                            Base.get_world_counter())
+    job = CompilerJob(source, config)
+
+    ir = sprint(io->GPUCompiler.code_llvm(io, job; dump_module=true))
+    @test occursin(r"load atomic ptr.* acquire", ir)
+    @test occursin(r"store atomic ptr.* release", ir)
+
+    air = sprint(io->GPUCompiler.code_native(io, job; dump_module=true))
+    @test !occursin(r"(load|store) atomic", air)
+    end
 end
 
 # byval lowering must strip the (non-IPO-safe) Julia const-region metadata off loads derived

test/spirv.jl

@@ -214,15 +214,17 @@ 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.
+@testset "atomic demotion" begin
+    # Julia's `unordered` heap-reference accesses and `release` type-tag stores 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_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)
+            y = Core.Intrinsics.atomic_pointerref(p, :acquire)
+            Core.Intrinsics.atomic_pointerset(q, y, :release)
             return
         end
     end