From 21b2053f0b5f379f0c6c8da1613e0d290e4befb9 Mon Sep 17 00:00:00 2001
From: Tim Besard <tim.besard@gmail.com>
Date: Thu, 24 Sep 2026 16:35:15 +0200
Subject: [PATCH] Metal: report unsupported atomics when lowering them.

`validate_ir` rejects the atomics that Metal cannot express, but when
validation is disabled (e.g., by `code_native` and the other reflection
macros), `lower_atomics!` selected them anyway, to AIR intrinsics that
don't exist (such as `air.atomic.global.store.i64`). Leave unsupported
operations alone instead, and report them after the lowering with the
same reasons and backtraces as `validate_ir`.

The fence lowering test relied on this, with a 64-bit atomic store.
---
 src/metal.jl  | 18 +++++++++++++++---
 test/metal.jl | 38 +++++++++++++++++++++++++++++++++++---
 2 files changed, 50 insertions(+), 6 deletions(-)

diff --git a/src/metal.jl b/src/metal.jl
index 5c9515c6d..09a092221 100644
--- a/src/metal.jl
+++ b/src/metal.jl
@@ -1437,8 +1437,7 @@ function lower_atomics!(@nospecialize(job::CompilerJob{MetalCompilerTarget}),
 
     for inst in atomics
         action = metal_atomic_action(job, inst)
-        # (`validate_ir` rejects unsupported operations; without validation, select them as-is)
-        action isa String && continue
+        action isa String && continue   # unsupported, reported below
         if action === :demote
             demote_private_atomic!(inst)
             continue
@@ -1461,10 +1460,23 @@ function lower_atomics!(@nospecialize(job::CompilerJob{MetalCompilerTarget}),
     end
 
     # select the atomics that are left, including the ones the expansions introduced
+    unsupported = IRError[]
     for f in functions(mod), bb in blocks(f), inst in collect(instructions(bb))
-        is_atomic_memop(inst) && select_atomic!(job, mod, inst)
+        is_atomic_memop(inst) || continue
+        action = metal_atomic_action(job, inst)
+        if action === :select
+            select_atomic!(job, mod, inst)
+        elseif action isa String
+            push!(unsupported, (action, backtrace(inst), string(inst)))
+        else
+            error("Atomic operation was not legalized ($action): $inst")
+        end
     end
 
+    # `validate_ir` rejects these, but validation can be disabled (e.g., for reflection).
+    # Rather than selecting AIR intrinsics that don't exist, fail here too.
+    isempty(unsupported) || throw(InvalidIRError(job, unsupported))
+
     # attach element type metadata to the declarations we introduced
     legalize_atomic_abi!(job, mod)
     return true
diff --git a/test/metal.jl b/test/metal.jl
index baf78780a..3386c97f7 100644
--- a/test/metal.jl
+++ b/test/metal.jl
@@ -2365,6 +2365,38 @@ end
                                            metal=v"3.1", air=v"2.6"))
     end
 
+    @testset "unsupported atomics without validation" begin
+        # the lowering reports them too, instead of selecting AIR intrinsics that don't exist
+        source = methodinstance(typeof(identity), Tuple{Int}, Base.get_world_counter())
+        target = MetalCompilerTarget(; macos=v"27", metal=v"4.1", air=v"2.9")
+        job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true))
+        Context(; opaque_pointers=true) do ctx
+            mod = parse(LLVM.Module, kernel("""
+                %a = atomicrmw add ptr addrspace(1) %p, i64 1 monotonic, align 8
+                %b = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4
+                """))
+            err = try
+                GPUCompiler.lower_atomics!(job, mod)
+                nothing
+            catch err
+                err
+            end
+            @test err isa GPUCompiler.InvalidIRError
+            @test startswith(only(err.errors)[1], "64-bit atomic operation")
+        end
+
+        # e.g. when reflecting, which compiles without validation
+        function kernel64(p::Core.LLVMPtr{Int64,1})
+            Core.Intrinsics.atomic_pointerset(reinterpret(Ptr{Int64}, p), 1, :monotonic)
+            return
+        end
+        source = methodinstance(typeof(kernel64), Tuple{Core.LLVMPtr{Int64,1}},
+                                Base.get_world_counter())
+        job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true))
+        @test occursin("store atomic i64", sprint(io -> GPUCompiler.code_llvm(io, job)))
+        @test_throws GPUCompiler.InvalidIRError GPUCompiler.code_native(devnull, job)
+    end
+
     # end-to-end, from Julia code emitting LLVM atomics (like UnsafeAtomics does), which also
     # exercises typed pointers on Julia versions that still use them
     mod = @eval module $(gensym())
@@ -2515,8 +2547,8 @@ end
 
     # end-to-end: Julia's `atomic_fence` intrinsic must not reach the AIR as a bare `fence`
     # (Julia 1.14 added a syncscope argument to the intrinsic, JuliaLang/julia#60311)
-    function kernel(p::Core.LLVMPtr{Int,1})
-        Core.Intrinsics.atomic_pointerset(reinterpret(Ptr{Int}, p), 1, :monotonic)
+    function kernel(p::Core.LLVMPtr{Int32,1})
+        Core.Intrinsics.atomic_pointerset(reinterpret(Ptr{Int32}, p), Int32(1), :monotonic)
         @static if VERSION >= v"1.14.0-DEV.1371"
             Core.Intrinsics.atomic_fence(:release, :system)
         else
@@ -2524,7 +2556,7 @@ end
         end
         return
     end
-    source = methodinstance(typeof(kernel), Tuple{Core.LLVMPtr{Int,1}}, Base.get_world_counter())
+    source = methodinstance(typeof(kernel), Tuple{Core.LLVMPtr{Int32,1}}, Base.get_world_counter())
     target = MetalCompilerTarget(; macos=v"27", metal=v"4.1", air=v"2.9")
     config = CompilerConfig(target, Metal.CompilerParams(); kernel=true)
     job = CompilerJob(source, config)