Skip to content

PTX: handle atomics on local memory #963

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
maleadt merged 2 commits into from
Oct 17, 2026
Merged

PTX: handle atomics on local memory #963

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c462734
PTX: demote atomics on thread-private memory.
maleadt Oct 16, 2026
e0b534d
PTX: check at run time whether generic atomics access local memory.
maleadt Oct 16, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
180 changes: 180 additions & 0 deletions src/ptx.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -455,6 +463,178 @@ 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, 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

end
return changed
end


## LLVM passes

# Replace occurrences of __nvvm_reflect("foo") and llvm.nvvm.reflect with an integer.
Expand Down
147 changes: 147 additions & 0 deletions test/ptx.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,6 +125,153 @@ 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 "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
Expand Down
Loading