From 849839a5f65802e60545f02f0ed9e7734af371a9 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 6 Oct 2026 15:57:19 +0200 Subject: [PATCH] Validate atomics for the PTX target. Atomics reach the PTX back-end from CUDA.jl's atomic functions, UnsafeAtomics and Atomix, Enzyme and Julia's intrinsics, so check them on the IR instead of in any one front-end. The NVPTX back-end doesn't reject everything the target can't run: before sm_60 it silently drops the system scope of read-modify-writes, it emits `.sys` 128-bit atomics that ptxas rejects before PTX ISA 8.4, and it aborts on unknown synchronization scopes and on sequentially-consistent 128-bit loads and stores. Other unsupported atomics (128-bit ones before sm_90, misaligned ones, atomics on local memory) fail in the back-end or in ptxas, without pointing at the Julia code that performs them. Platforms can also reject system-scope atomics the ISA supports, e.g. Pascal GPUs under Windows. Front-ends report that through the new `system_atomics` target field. --- src/ptx.jl | 110 +++++++++++++++++++++++++++++++++++++++++++ test/helpers/ptx.jl | 4 +- test/ptx.jl | 111 ++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 223 insertions(+), 2 deletions(-) diff --git a/src/ptx.jl b/src/ptx.jl index f97bddc78..d36b0038c 100644 --- a/src/ptx.jl +++ b/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. diff --git a/test/helpers/ptx.jl b/test/helpers/ptx.jl index 911275440..612e78db9 100644 --- a/test/helpers/ptx.jl +++ b/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 diff --git a/test/ptx.jl b/test/ptx.jl index b6319d263..6264673bf 100644 --- a/test/ptx.jl +++ b/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)