diff --git a/src/ptx.jl b/src/ptx.jl index d36b0038..2b913440 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,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. diff --git a/test/ptx.jl b/test/ptx.jl index 6264673b..3ab458ce 100644 --- a/test/ptx.jl +++ b/test/ptx.jl @@ -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