From c462734b8e2b6618c90af4aecf3d900dc17ec9b6 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 16 Oct 2026 16:45:31 +0200 Subject: [PATCH 1/2] PTX: demote atomics on thread-private memory. PTX has no atomics on the stack (local memory): ptxas rejects the `atom.local.cas` that NVPTX selects for a `cmpxchg` on an `alloca`, atomic loads and stores on the stack become `ld.acquire` and `st.release` of generic addresses, which PTX leaves undefined, and NVPTX only demotes `atomicrmw` on pointers it knows to be local, and only when optimizing. Atomics on stack slots also failed validation for no reason, e.g., system-scope ones before sm_60. Like on Metal and SPIR-V, replace atomics on the thread's own memory by plain accesses from `finish_ir!`, after optimization and before validation. On PTX, pointers in the local address space are private too. Co-authored-by: Tim Esselens --- src/ptx.jl | 43 +++++++++++++++++++++++++++++++++++++++++++ test/ptx.jl | 49 +++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 92 insertions(+) diff --git a/src/ptx.jl b/src/ptx.jl index d36b0038..25ddda3c 100644 --- a/src/ptx.jl +++ b/src/ptx.jl @@ -303,6 +303,14 @@ function optimize_module!(@nospecialize(job::CompilerJob{PTXCompilerTarget}), end end +function finish_ir!(@nospecialize(job::CompilerJob{PTXCompilerTarget}), mod::LLVM.Module, + entry::LLVM.Function) + # PTX has no atomics on the stack, and NVPTX doesn't legalize them + lower_local_atomics!(mod) + + return entry +end + # mirrors `NVPTXCompileOptions` from libnvptx.h struct NVPTXCompileOptions cpu::Cstring @@ -455,6 +463,41 @@ function nvptx_atomic_error(inst::LLVM.Instruction) end +## atomics on local memory + +# PTX only has atomic instructions, and atomic (relaxed, acquire or release) loads and +# stores, for global and shared memory, and for generic addresses of those. On the stack +# (local memory) they are undefined: Pascal GPUs fail with CUDA_ERROR_INVALID_ADDRESS_SPACE, +# which leaves the context unusable, and ptxas rejects the `atom.local` that NVPTX selects +# for a `cmpxchg` on an `alloca`. NVPTX only demotes `atomicrmw` on pointers it knows to be +# local, and only when optimizing. +# +# Atomics on the stack are common nonetheless, e.g., on objects that Julia's `AllocOpt` +# moved there. Since no other thread can access a thread's stack, `lower_local_atomics!` +# replaces them by plain accesses. + +# Is `ptr` known to point to the thread's own memory? +ptx_thread_private(ptr::LLVM.Value) = + ptr.value_type.addrspace == #=local=# 5 || is_thread_private(ptr) + +function lower_local_atomics!(mod::LLVM.Module) + changed = false + @tracepoint "lower local atomics" begin + + atomics = [inst for f in mod.functions for bb in f.blocks for inst in bb.instructions + if is_atomic_memop(inst)] + for inst in atomics + if ptx_thread_private(inst.pointer_operand) + demote_private_atomic!(inst) + changed = true + end + end + + end + return changed +end + + ## LLVM passes # Replace occurrences of __nvvm_reflect("foo") and llvm.nvvm.reflect with an integer. diff --git a/test/ptx.jl b/test/ptx.jl index 6264673b..74766e6c 100644 --- a/test/ptx.jl +++ b/test/ptx.jl @@ -125,6 +125,55 @@ end hash(PTXCompilerTarget(cap=v"6.1", system_atomics=false)) end +@testset "atomics on the stack" begin + # atomics on the thread's own memory become plain accesses + # (typed-pointer IR, which also parses with opaque pointers) + mod = @eval module $(gensym()) + rmw(v) = Base.llvmcall(""" + %a = alloca i32, align 4 + store i32 1, i32* %a, align 4 + %old = atomicrmw add i32* %a, i32 %0 monotonic, align 4 + ret i32 %old""", Int32, Tuple{Int32}, v) + cmpxchg(v) = Base.llvmcall(""" + %a = alloca i32, align 4 + store i32 1, i32* %a, align 4 + %r = cmpxchg volatile i32* %a, i32 1, i32 %0 seq_cst seq_cst, align 4 + %old = extractvalue { i32, i1 } %r, 0 + ret i32 %old""", Int32, Tuple{Int32}, v) + load_store(v) = Base.llvmcall(""" + %a = alloca i32, align 4 + %g = addrspacecast i32* %a to i32 addrspace(5)* + store atomic i32 %0, i32 addrspace(5)* %g release, align 4 + %old = load atomic i32, i32* %a acquire, align 4 + ret i32 %old""", Int32, Tuple{Int32}, v) + local_arg(p, v) = Base.llvmcall(""" + %p = bitcast i8 addrspace(5)* %0 to i32 addrspace(5)* + %old = atomicrmw xchg i32 addrspace(5)* %p, i32 %1 seq_cst, align 4 + ret i32 %old""", Int32, Tuple{Core.LLVMPtr{Int32,5}, Int32}, p, v) + kernel(v) = (rmw(v); nothing) + end + for (f, types) in ((mod.rmw, Tuple{Int32}), (mod.cmpxchg, Tuple{Int32}), + (mod.load_store, Tuple{Int32}), + (mod.local_arg, Tuple{Core.LLVMPtr{Int32,5}, Int32})) + for optimize in (true, false) + ir = sprint(io->PTX.code_llvm(io, f, types; optimize)) + @test !occursin(r"atomicrmw |cmpxchg |load atomic |store atomic ", ir) + end + # NVPTX selects an `atom.local` for `cmpxchg`, which ptxas rejects + asm = sprint(io->PTX.code_native(io, f, types)) + @test !occursin(r"\batom\.|\.(relaxed|acquire|release)\.", asm) + end + # the volatile `cmpxchg` keeps its volatile accesses + @test @filecheck begin + @check "load volatile i32" + @check "store volatile i32" + PTX.code_llvm(mod.cmpxchg, Tuple{Int32}) + end + + # they aren't validated either, as they don't need atomic instructions + @test PTX.code_execution(mod.kernel, Tuple{Int32}; cap=v"5.2") isa Tuple +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 From e0b534de612b10a2d86dd1fad3b9a3d2b1831f40 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 16 Oct 2026 16:53:16 +0200 Subject: [PATCH 2/2] PTX: check at run time whether generic atomics access local memory. Atomics on generic pointers that point to the stack are undefined in PTX, and fail on Pascal GPUs with CUDA_ERROR_INVALID_ADDRESS_SPACE, which leaves the context unusable. Enzyme emits them when it accumulates into the shadow of an argument without seeing where the argument points, e.g., when a stack slot of the caller is passed to a function that isn't inlined (Enzyme.jl#428). Inlining hides the problem only sometimes. Like NVIDIA's libcu++ does for `cuda::atomic_ref`, check `isspacep.local` before atomic operations on generic pointers that may point to the stack, and perform the operation with plain accesses on local memory. The atomic that is kept for other memory gets `!noalias.addrspace` metadata, so that it isn't checked again. Pointers known to point to global or shared memory are left alone, and NVPTX folds the check when it can infer the address space, e.g., for pointers passed to a kernel. Julia's `unordered` accesses of heap references aren't checked: they are not used to synchronize, and NVPTX only lowers them to the relaxed loads and stores that local memory doesn't support from sm_70 on, where they haven't been seen to fail. Co-authored-by: Tim Esselens --- src/ptx.jl | 141 +++++++++++++++++++++++++++++++++++++++++++++++++++- test/ptx.jl | 98 ++++++++++++++++++++++++++++++++++++ 2 files changed, 237 insertions(+), 2 deletions(-) diff --git a/src/ptx.jl b/src/ptx.jl index 25ddda3c..2b913440 100644 --- a/src/ptx.jl +++ b/src/ptx.jl @@ -473,23 +473,160 @@ end # local, and only when optimizing. # # Atomics on the stack are common nonetheless, e.g., on objects that Julia's `AllocOpt` -# moved there. Since no other thread can access a thread's stack, `lower_local_atomics!` -# replaces them by plain accesses. +# moved there, or in code from Enzyme, which accumulates into the shadow of an argument with +# an atomic when it cannot see where the argument points, e.g., to a stack slot of the +# caller (Enzyme.jl#428). Since no other thread can access a thread's stack, plain accesses +# behave the same, so `lower_local_atomics!`, like NVIDIA's libcu++: +# - replaces atomics on the thread's own memory by plain accesses; +# - checks at run time whether an atomic on a generic pointer that may point to the stack +# does, and then performs the plain accesses instead (see `guard_local_atomic!`). +# +# `unordered` loads and stores aren't checked: Julia uses them for every access to a heap +# reference, and NVPTX only lowers them to instructions that local memory doesn't support +# (`ld.relaxed` and `st.relaxed`) from sm_70 on, which hasn't been seen to fail. # Is `ptr` known to point to the thread's own memory? ptx_thread_private(ptr::LLVM.Value) = ptr.value_type.addrspace == #=local=# 5 || is_thread_private(ptr) +# Can `ptr` point to local memory? Generic pointers can, unless every object they can be +# derived from is in another address space. +function ptx_may_be_local(ptr::LLVM.Value) + seen = Set{LLVM.Value}() + worklist = LLVM.Value[ptr] + while !isempty(worklist) + val = pop!(worklist) + val in seen && continue + push!(seen, val) + as = val.value_type.addrspace + if as == #=local=# 5 + return true + elseif as != #=generic=# 0 || val isa LLVM.GlobalVariable + continue + elseif val isa LLVM.GetElementPtrInst || val isa LLVM.BitCastInst || + val isa LLVM.AddrSpaceCastInst + push!(worklist, val.operands[1]) + elseif val isa LLVM.ConstantExpr && + LLVM.opcode(val) in (LLVM.API.LLVMGetElementPtr, LLVM.API.LLVMBitCast, + LLVM.API.LLVMAddrSpaceCast) + push!(worklist, val.operands[1]) + elseif val isa LLVM.PHIInst + append!(worklist, first.(val.incoming)) + elseif val isa LLVM.SelectInst + push!(worklist, val.operands[2], val.operands[3]) + else + return true + end + end + return false +end + +# Does the `!noalias.addrspace` metadata of `inst` rule out that it accesses local memory? +# `guard_local_atomic!` adds it to the atomic it keeps for other memory. +function excludes_local(inst::LLVM.Instruction) + haskey(inst.metadata, "noalias.addrspace") || return false + md = inst.metadata["noalias.addrspace"] + bounds = [convert(Int, Value(op)) % UInt32 for op in md.operands] + # (pairs of half-open ranges, which can wrap around) + return any(1:2:length(bounds)) do i + lo, hi = bounds[i], bounds[i+1] + lo <= hi ? lo <= 5 < hi : (5 >= lo || 5 < hi) + end +end + +# Replace the atomic memory operation `inst` by a call to a function that checks whether +# its pointer points to local memory, and then performs it with plain accesses, or atomically +# otherwise. The function is meant to be inlined (it is `alwaysinline`). +function guard_local_atomic!(inst::LLVM.Instruction) + mod = LLVM.parent(LLVM.parent(LLVM.parent(inst))) + args = collect(inst.operands) + T = inst.value_type + ft = LLVM.FunctionType(T, [arg.value_type for arg in args]) + f = LLVM.Function(mod, "gpucompiler.local_atomic", ft) + f.linkage = LLVM.Linkage.Internal + push!(f.function_attributes, EnumAttribute(:alwaysinline)) + + # the operation, on the arguments of the function. the clones would refer to the debug + # info of the caller, so they get the location of the call when inlined instead. + local_atomic, other_atomic = copy(inst), copy(inst) + for clone in (local_atomic, other_atomic) + clone.debug_location = nothing + for (i, param) in enumerate(f.parameters) + clone.operands[i] = param + end + end + other_atomic.metadata["noalias.addrspace"] = + MDNode([Metadata(ConstantInt(Int32(5))), Metadata(ConstantInt(Int32(6)))]) + + isspacep = LLVM.Function(mod, LLVM.Intrinsic("llvm.nvvm.isspacep.local")) + isspacep_ft = isspacep.function_type + @dispose builder=IRBuilder() begin + entry = BasicBlock(f, "entry") + local_block = BasicBlock(f, "local") + other_block = BasicBlock(f, "other") + exit = BasicBlock(f, "exit") + + position!(builder, LLVM.at_end(entry)) + # (the intrinsic takes an `i8*` with typed pointers) + ptr = f.parameters[inst isa LLVM.StoreInst ? 2 : 1] + T_ptr = only(isspacep_ft.parameters) + ptr = ptr.value_type == T_ptr ? ptr : bitcast!(builder, ptr, T_ptr) + is_local = call!(builder, isspacep_ft, isspacep, [ptr]) + br!(builder, is_local, local_block, other_block) + + for (block, atomic) in ((local_block, local_atomic), (other_block, other_atomic)) + position!(builder, LLVM.at_end(block)) + LLVM.API.LLVMInsertIntoBuilder(builder, atomic) + br!(builder, exit) + end + + position!(builder, LLVM.at_end(exit)) + if T isa LLVM.VoidType + ret!(builder) + else + result = phi!(builder, T) + append!(result.incoming, [(local_atomic, local_block), (other_atomic, other_block)]) + ret!(builder, result) + end + end + # (after building the users of its result, which lowering replaces) + demote_private_atomic!(local_atomic) + + @dispose builder=IRBuilder() begin + position!(builder, LLVM.before(inst)) + call = call!(builder, ft, f, args) + call.debug_location = inst.debug_location + T isa LLVM.VoidType || replace_uses!(inst, call) + end + erase!(inst) + return +end + function lower_local_atomics!(mod::LLVM.Module) changed = false @tracepoint "lower local atomics" begin atomics = [inst for f in mod.functions for bb in f.blocks for inst in bb.instructions if is_atomic_memop(inst)] + guarded = false for inst in atomics if ptx_thread_private(inst.pointer_operand) demote_private_atomic!(inst) changed = true + elseif (inst isa LLVM.LoadInst || inst isa LLVM.StoreInst) && + inst.ordering == LLVM.AtomicOrdering.Unordered + continue + elseif ptx_may_be_local(inst.pointer_operand) && !excludes_local(inst) + guard_local_atomic!(inst) + changed = guarded = true + end + end + + # inline the checks (the pass removes the functions) + if guarded + @dispose pb=PassBuilder() begin + add!(pb, AlwaysInlinerPass()) + run!(pb, mod) end end diff --git a/test/ptx.jl b/test/ptx.jl index 74766e6c..3ab458ce 100644 --- a/test/ptx.jl +++ b/test/ptx.jl @@ -174,6 +174,104 @@ end @test PTX.code_execution(mod.kernel, Tuple{Int32}; cap=v"5.2") isa Tuple end +@testset "atomics on generic pointers" begin + # atomics on generic pointers that may point to the stack check whether they do + mod = @eval module $(gensym()) + rmw(p, v) = Base.llvmcall(""" + %p = inttoptr i64 %0 to i32* + %old = atomicrmw add i32* %p, i32 %1 syncscope("block") monotonic, align 4 + ret i32 %old""", Int32, Tuple{Ptr{Int32}, Int32}, p, v) + cmpxchg(p, v) = Base.llvmcall(""" + %p = inttoptr i64 %0 to i32* + %r = cmpxchg weak volatile i32* %p, i32 1, i32 %1 seq_cst acquire, align 4 + %old = extractvalue { i32, i1 } %r, 0 + ret i32 %old""", Int32, Tuple{Ptr{Int32}, Int32}, p, v) + load_store(p, v) = Base.llvmcall(""" + %p = inttoptr i64 %0 to i32* + store atomic i32 %1, i32* %p release, align 4 + %old = load atomic i32, i32* %p monotonic, align 4 + ret i32 %old""", Int32, Tuple{Ptr{Int32}, Int32}, p, v) + unordered(p) = Base.llvmcall(""" + %p = inttoptr i64 %0 to i32* + %old = load atomic i32, i32* %p unordered, align 4 + ret i32 %old""", Int32, Tuple{Ptr{Int32}}, p) + global_cast(p, v) = Base.llvmcall(""" + %p = bitcast i8 addrspace(1)* %0 to i32 addrspace(1)* + %g = addrspacecast i32 addrspace(1)* %p to i32* + %old = atomicrmw add i32* %g, i32 %1 monotonic, align 4 + ret i32 %old""", Int32, Tuple{Core.LLVMPtr{Int32,1}, Int32}, p, v) + end + + for optimize in (true, false) + # the original operation for other memory, plain accesses for local memory + @test @filecheck begin + @check "call i1 @llvm.nvvm.isspacep.local" + @check "[[OLD:%[^ ]+]] = load i32," + @check "[[NEW:%[^ ]+]] = add i32 [[OLD]]" + @check "store i32 [[NEW]]" + @check "atomicrmw add {{.*}} syncscope(\"block\") monotonic, align 4{{.*}}, !noalias.addrspace" + @check "phi i32 [ [[OLD]]" + @check_not "atomicrmw" + PTX.code_llvm(mod.rmw, Tuple{Ptr{Int32}, Int32}; optimize, raw=true) + end + @test @filecheck begin + @check "call i1 @llvm.nvvm.isspacep.local" + @check "load volatile i32" + @check "store volatile i32" + @check "cmpxchg weak volatile {{.*}} seq_cst acquire, align 4{{.*}}, !noalias.addrspace" + PTX.code_llvm(mod.cmpxchg, Tuple{Ptr{Int32}, Int32}; optimize, raw=true) + end + @test @filecheck begin + @check "call i1 @llvm.nvvm.isspacep.local" + @check "store atomic i32 {{.*}} release, align 4{{.*}}, !noalias.addrspace" + @check "call i1 @llvm.nvvm.isspacep.local" + @check "load atomic i32{{.*}} monotonic, align 4{{.*}}, !noalias.addrspace" + PTX.code_llvm(mod.load_store, Tuple{Ptr{Int32}, Int32}; optimize, raw=true) + end + end + + # not for Julia's unordered heap-reference accesses, or pointers to other memory + @test @filecheck begin + @check_not "isspacep" + @check "load atomic i32{{.*}} unordered" + PTX.code_llvm(mod.unordered, Tuple{Ptr{Int32}}) + end + @test @filecheck begin + @check_not "isspacep" + @check "atomicrmw add" + PTX.code_llvm(mod.global_cast, Tuple{Core.LLVMPtr{Int32,1}, Int32}) + end + + # checked atomics aren't checked again, and the checks are inlined + Context() do ctx + ir = parse(LLVM.Module, """ + define i32 @f(i32* %p, i32 %v) { + %old = atomicrmw add i32* %p, i32 %v monotonic, align 4 + ret i32 %old + }""") + @test GPUCompiler.lower_local_atomics!(ir) + @test !GPUCompiler.lower_local_atomics!(ir) + @test [LLVM.name(f) for f in ir.functions] == ["f", "llvm.nvvm.isspacep.local"] + end + # neither are ones whose `!noalias.addrspace` otherwise rules out local memory + Context() do ctx + for (ranges, checked) in (("i32 5, i32 6", false), ("i32 0, i32 -1", false), + ("i32 7, i32 6", false), ("i32 3, i32 4", true)) + ir = parse(LLVM.Module, """ + define void @f(i32* %p) { + %old = atomicrmw add i32* %p, i32 1 monotonic, align 4, !noalias.addrspace !0 + ret void + } + !0 = !{$ranges}""") + @test GPUCompiler.lower_local_atomics!(ir) == checked + end + end + + asm = sprint(io->PTX.code_native(io, mod.cmpxchg, Tuple{Ptr{Int32}, Int32})) + @test occursin("isspacep.local", asm) + @test occursin("atom.acquire.sys.cas.b32", asm) +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