Validate atomics for the PTX target

src/ptx.jl

@@ -31,6 +31,11 @@ Base.@kwdef struct PTXCompilerTarget <: AbstractCompilerTarget
 
     fastmath::Bool = Base.JLOptions().fast_math == 1
 
+    # can system-scope atomics be used? the ISA supports them from sm_60 on (which
+    # `validate_ir` checks separately), but some platforms reject them, e.g. Pascal GPUs
+    # under Windows, or Tegra GPUs before sm_72.
+    system_atomics::Bool = true
+
     # deprecated; remove with next major version
     exitable::Union{Nothing,Bool} = nothing
     unreachable::Union{Nothing,Bool} = nothing
@@ -48,6 +53,7 @@ function Base.hash(target::PTXCompilerTarget, h::UInt)
     h = hash(target.blocks_per_sm, h)
     h = hash(target.maxregs, h)
     h = hash(target.fastmath, h)
+    h = hash(target.system_atomics, h)
 
     h
 end
@@ -345,6 +351,110 @@ function llvm_debug_info(@nospecialize(job::CompilerJob{PTXCompilerTarget}))
 end
 
 
+## validation
+
+# Atomic operations and fences come from many front-ends (CUDA.jl's atomic functions,
+# UnsafeAtomics and Atomix, Enzyme, Julia's atomic intrinsics), so they are validated here,
+# on the IR, rather than in any of them. The NVPTX back-end does not reject everything the
+# target cannot run: it silently drops a system scope before sm_60, emits PTX that `ptxas`
+# or the driver reject, or aborts the process. Validation runs after `lower_syncscopes!`, so
+# the scopes are the ones NVPTX knows.
+function validate_ir(job::CompilerJob{PTXCompilerTarget}, mod::LLVM.Module)
+    errors = IRError[]
+    dl = mod.datalayout
+    for f in mod.functions, bb in f.blocks, inst in bb.instructions
+        reason = if is_atomic_memop(inst)
+            ptx_atomic_error(job, dl, inst)
+        elseif inst isa LLVM.FenceInst
+            ptx_syncscope_error(inst, "fence")
+        else
+            nothing
+        end
+        reason === nothing || push!(errors, (reason, backtrace(inst), string(inst)))
+    end
+    return errors
+end
+
+# the synchronization scopes of NVPTX, as `lower_syncscopes!` leaves them. a cluster is a
+# single block on targets without clusters, where NVPTX uses the block scope instead.
+const PTX_SYNCSCOPES = ("singlethread", "block", "cluster", "device", "system")
+
+ptx_syncscope_error(inst::LLVM.Instruction, what="atomic operation") =
+    inst.syncscope.name in PTX_SYNCSCOPES ? nothing :
+    "$what with synchronization scope $(repr(inst.syncscope.name))"
+
+# Why the target cannot run the atomic memory operation `inst`, or `nothing` if it can.
+# Operations PTX has no instruction for (e.g. 8- and 16-bit read-modify-writes, most
+# floating-point ones, or `fadd` on older devices) are fine: NVPTX expands them to
+# compare-exchange loops.
+function ptx_atomic_error(@nospecialize(job::CompilerJob{PTXCompilerTarget}),
+                          dl::LLVM.DataLayout, inst::LLVM.Instruction)
+    target = job.config.target
+    reason = ptx_syncscope_error(inst)
+    reason === nothing || return reason
+    system = inst.syncscope.name == "system"
+
+    as = inst.pointer_operand.value_type.addrspace
+    if !(as in (0, 1, 3, 7))
+        return "atomic operation in address space $as (PTX only supports atomics on generic, global and shared memory)"
+    end
+    if as == 7 && (target.cap < v"9.0" || target.ptx < v"7.8")
+        return "atomic operation on shared::cluster memory (requires compute capability 9.0 and PTX ISA 7.8)"
+    end
+
+    # (vectors aren't supported either: integer vectors crash NVPTX 23.1, and PTX only has
+    # vector atomics for a few floating-point reductions)
+    T = atomic_value_type(inst)
+    bits = if T isa LLVM.IntegerType || T isa LLVM.FloatingPointType || T isa LLVM.PointerType
+        Int(LLVM.bit_size(dl, T))
+    end
+    if !(bits in (8, 16, 32, 64) || (bits == 128 && T isa LLVM.IntegerType))
+        return "atomic operation on a $(string(T)) value"
+    end
+    if bits == 128
+        if target.cap < v"9.0" || target.ptx < v"8.3"
+            return "128-bit atomic operation (requires compute capability 9.0 and PTX ISA 8.3)"
+        end
+        # (NVPTX emits `.sys` 128-bit atomics with PTX ISA 8.3, which `ptxas` rejects)
+        if system && target.ptx < v"8.4"
+            return "system-scope 128-bit atomic operation (requires PTX ISA 8.4)"
+        end
+    end
+    if inst.alignment < bits ÷ 8
+        return "atomic operation with alignment $(inst.alignment) (requires at least $(bits ÷ 8)-byte alignment)"
+    end
+
+    # loads and stores only take a scope from sm_70 on (before, they are volatile accesses,
+    # bracketed by `membar`s that support every scope), and shared memory is only visible
+    # within a block or cluster, so this only concerns read-modify-writes on other memory
+    if system && (inst isa LLVM.AtomicRMWInst || inst isa LLVM.AtomicCmpXchgInst) &&
+       !(as in (3, 7))
+        if target.cap < v"6.0"
+            return "system-scope atomic operation (requires compute capability 6.0; use device scope if system-wide atomicity is not required)"
+        elseif !target.system_atomics
+            return "system-scope atomic operation (not supported on this platform; use device scope if system-wide atomicity is not required)"
+        end
+    end
+
+    return nvptx_atomic_error(inst)
+end
+
+# Why the NVPTX back-end cannot compile the atomic memory operation `inst`, although the
+# target supports it, or `nothing` if it can. These work around bugs in the back-end, and
+# should be revisited when updating it.
+function nvptx_atomic_error(inst::LLVM.Instruction)
+    # NVPTX's instruction printer aborts on the ordering of the `atom.cas.b128` and
+    # `atom.exch.b128` that it legalizes these to (LLVM 23.1, and main as of October 2026)
+    T = atomic_value_type(inst)
+    if (inst isa LLVM.LoadInst || inst isa LLVM.StoreInst) &&
+       inst.ordering == LLVM.AtomicOrdering.SequentiallyConsistent &&
+       T isa LLVM.IntegerType && T.width == 128
+        return "sequentially-consistent 128-bit atomic load or store (not supported by the NVPTX back-end)"
+    end
+    return nothing
+end
+
+
 ## LLVM passes
 
 # Replace occurrences of __nvvm_reflect("foo") and llvm.nvvm.reflect with an integer.

test/helpers/ptx.jl

@@ -47,13 +47,13 @@ function create_job(@nospecialize(func), @nospecialize(types);
                     cap=v"7.0", ptx=v"6.0", feature_set=:baseline,
                     minthreads=nothing, maxthreads=nothing,
                     blocks_per_sm=nothing, maxregs=nothing,
-                    fastmath=false, patch::Bool=false,
+                    fastmath=false, system_atomics=true, patch::Bool=false,
                     kwargs...)
     config_kwargs, kwargs = split_kwargs(kwargs, GPUCompiler.CONFIG_KWARGS)
     source = methodinstance(typeof(func), Base.to_tuple_type(types), Base.get_world_counter())
     target = PTXCompilerTarget(; cap, ptx, feature_set,
                                  minthreads, maxthreads, blocks_per_sm, maxregs,
-                                 fastmath)
+                                 fastmath, system_atomics)
     params = CompilerParams(patch)
     config = CompilerConfig(target, params; kernel=false, config_kwargs...)
     CompilerJob(source, config), kwargs

test/ptx.jl

@@ -40,6 +40,91 @@ end
     end
 end
 
+@testset "atomic validation" begin
+    # validate textual IR, with NVPTX's synchronization scopes, for a PTX target
+    function validate_atomics(body; cap, ptx=v"8.4", system_atomics=true)
+        source = methodinstance(typeof(identity), Tuple{Int}, Base.get_world_counter())
+        target = PTXCompilerTarget(; cap, ptx, system_atomics)
+        job = CompilerJob(source, CompilerConfig(target, PTX.CompilerParams(); kernel=true))
+        Context(; opaque_pointers=true) do ctx
+            mod = parse(LLVM.Module, """
+                target datalayout = "$(GPUCompiler.llvm_datalayout(target))"
+                define void @kernel(ptr addrspace(1) %p, ptr addrspace(3) %s,
+                                    ptr addrspace(7) %c, ptr addrspace(5) %l, ptr %g) {
+                  $body
+                  ret void
+                }""")
+            join(first.(GPUCompiler.validate_ir(job, mod)), "\n")
+        end
+    end
+
+    @testset "unsupported" begin
+        for (body, target, reason) in (
+                ("%a = atomicrmw add ptr addrspace(1) %p, i32 1 syncscope(\"agent\") monotonic, align 4",
+                 (; cap=v"7.0"), "atomic operation with synchronization scope \"agent\""),
+                ("fence syncscope(\"agent\") seq_cst",
+                 (; cap=v"7.0"), "fence with synchronization scope \"agent\""),
+                ("%a = atomicrmw add ptr addrspace(5) %l, i32 1 syncscope(\"device\") monotonic, align 4",
+                 (; cap=v"7.0"), "atomic operation in address space 5"),
+                ("%a = atomicrmw add ptr addrspace(7) %c, i32 1 syncscope(\"cluster\") monotonic, align 4",
+                 (; cap=v"8.0"), "atomic operation on shared::cluster memory"),
+                ("%a = load atomic fp128, ptr addrspace(1) %p syncscope(\"device\") monotonic, align 16",
+                 (; cap=v"9.0"), "atomic operation on a fp128 value"),
+                ("%a = cmpxchg ptr addrspace(1) %p, i128 0, i128 1 syncscope(\"device\") monotonic monotonic, align 16",
+                 (; cap=v"8.0", ptx=v"7.8"), "128-bit atomic operation (requires compute capability 9.0 and PTX ISA 8.3)"),
+                ("%a = atomicrmw xchg ptr addrspace(1) %p, i128 1 syncscope(\"device\") monotonic, align 16",
+                 (; cap=v"9.0", ptx=v"8.0"), "128-bit atomic operation (requires compute capability 9.0 and PTX ISA 8.3)"),
+                ("%a = cmpxchg ptr addrspace(1) %p, i128 0, i128 1 monotonic monotonic, align 16",
+                 (; cap=v"9.0", ptx=v"8.3"), "system-scope 128-bit atomic operation (requires PTX ISA 8.4)"),
+                ("%a = atomicrmw add ptr addrspace(1) %p, i64 1 syncscope(\"device\") monotonic, align 4",
+                 (; cap=v"7.0"), "atomic operation with alignment 4 (requires at least 8-byte alignment)"),
+                ("%a = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4",
+                 (; cap=v"5.2"), "system-scope atomic operation (requires compute capability 6.0"),
+                ("%a = cmpxchg ptr %g, i32 0, i32 1 monotonic monotonic, align 4",
+                 (; cap=v"5.2"), "system-scope atomic operation (requires compute capability 6.0"),
+                ("%a = atomicrmw fadd ptr addrspace(1) %p, double 1.0 seq_cst, align 8",
+                 (; cap=v"5.2"), "system-scope atomic operation (requires compute capability 6.0"),
+                ("%a = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4",
+                 (; cap=v"6.1", system_atomics=false), "system-scope atomic operation (not supported on this platform"),
+                ("%a = load atomic i128, ptr addrspace(1) %p syncscope(\"device\") seq_cst, align 16",
+                 (; cap=v"9.0"), "sequentially-consistent 128-bit atomic load or store"),
+                ("store atomic i128 0, ptr addrspace(1) %p syncscope(\"device\") seq_cst, align 16",
+                 (; cap=v"9.0"), "sequentially-consistent 128-bit atomic load or store"),
+            )
+            @test occursin(reason, validate_atomics(body; target...))
+        end
+    end
+
+    @testset "supported" begin
+        for (body, target) in (
+                # system scope where it doesn't need system-scope atomics
+                ("%a = atomicrmw add ptr addrspace(3) %s, i32 1 monotonic, align 4", (; cap=v"5.2")),
+                ("%a = load atomic i32, ptr addrspace(1) %p acquire, align 4\nstore atomic i32 %a, ptr addrspace(1) %p release, align 4",
+                 (; cap=v"5.2", system_atomics=false)),
+                ("fence seq_cst", (; cap=v"5.2", system_atomics=false)),
+                ("%a = atomicrmw add ptr addrspace(1) %p, i32 1 syncscope(\"device\") monotonic, align 4", (; cap=v"5.2")),
+                ("%a = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4", (; cap=v"6.0")),
+                # a cluster is a block before sm_90
+                ("%a = atomicrmw add ptr addrspace(1) %p, i32 1 syncscope(\"cluster\") monotonic, align 4\nfence syncscope(\"cluster\") acquire",
+                 (; cap=v"8.0", ptx=v"7.8")),
+                ("%a = atomicrmw add ptr addrspace(7) %c, i32 1 syncscope(\"cluster\") monotonic, align 4", (; cap=v"9.0", ptx=v"7.8")),
+                # operations NVPTX expands
+                ("%a = atomicrmw nand ptr addrspace(1) %p, i8 1 syncscope(\"block\") monotonic, align 1", (; cap=v"5.2")),
+                ("%a = atomicrmw fmax ptr %g, half 1.0 syncscope(\"singlethread\") monotonic, align 2", (; cap=v"5.2")),
+                ("%a = load atomic ptr, ptr addrspace(1) %p unordered, align 8", (; cap=v"5.2")),
+                ("%a = cmpxchg ptr addrspace(1) %p, i128 0, i128 1 syncscope(\"device\") seq_cst seq_cst, align 16",
+                 (; cap=v"9.0", ptx=v"8.3")),
+                ("%a = load atomic i128, ptr addrspace(1) %p acquire, align 16", (; cap=v"9.0", ptx=v"8.4")),
+            )
+            @test validate_atomics(body; target...) == ""
+        end
+    end
+
+    # the platform restriction is part of the target identity
+    @test hash(PTXCompilerTarget(cap=v"6.1")) !=
+          hash(PTXCompilerTarget(cap=v"6.1", system_atomics=false))
+end
+
 @testset "kernel state survives a runtime rebuild" begin
     # Clearing the runtime cache forces the library link inside `emit_llvm` to rebuild
     # the runtime (nested compilation); the kernel must still get its state argument
@@ -544,6 +629,32 @@ end
     end
 end
 
+@testset "atomic validation" begin
+    # atomics in Julia code are reported with the frame that performs them.
+    # emit the atomic directly, as Julia's own intrinsics lower to a CAS loop calling `+`
+    # (dynamically, on 1.10) or fold an unused swap into a store, depending on the version.
+    mod = @eval module $(gensym())
+        function kernel(p::Ptr{Int})
+            Base.llvmcall("""%ptr = inttoptr i64 %0 to i64*
+                             %old = atomicrmw add i64* %ptr, i64 1 monotonic, align 8
+                             ret void""", Cvoid, Tuple{Ptr{Int}}, p)
+            return
+        end
+    end
+    @test_throws_message(InvalidIRError,
+                         PTX.code_execution(mod.kernel, Tuple{Ptr{Int}}; cap=v"5.2")) do msg
+        occursin("Reason: unsupported system-scope atomic operation (requires compute capability 6.0", msg) &&
+        occursin(r"\[\d+\] kernel", msg)
+    end
+    @test_throws_message(InvalidIRError,
+                         PTX.code_execution(mod.kernel, Tuple{Ptr{Int}}; cap=v"6.1",
+                                            system_atomics=false)) do msg
+        occursin("Reason: unsupported system-scope atomic operation (not supported on this platform", msg)
+    end
+    asm, _ = PTX.code_execution(mod.kernel, Tuple{Ptr{Int}}; cap=v"6.1")
+    @test occursin("atom.sys.global.add.u64", asm)
+end
+
 @testset "float boxes" begin
     mod = @eval module $(gensym())
         function kernel(a,b)