diff --git a/src/GPUCompiler.jl b/src/GPUCompiler.jl index b74253ca..5599e2d1 100644 --- a/src/GPUCompiler.jl +++ b/src/GPUCompiler.jl @@ -72,6 +72,7 @@ include("mangling.jl") # compiler interface and implementations include("interface.jl") include("relocation.jl") +include("atomics.jl") include("error.jl") include("native.jl") include("ptx.jl") diff --git a/src/atomics.jl b/src/atomics.jl new file mode 100644 index 00000000..984f14c6 --- /dev/null +++ b/src/atomics.jl @@ -0,0 +1,110 @@ +# target-independent helpers for lowering LLVM atomics +# +# Back-ends without an LLVM target (Metal) or with one that mishandles some atomics (SPIR-V) +# need to inspect and rewrite LLVM's atomic instructions themselves; these helpers mirror the +# parts of LLVM's `AtomicExpand` and the atomic instruction classes that LLVM.jl lacks. + +const AtomicOrdering = LLVM.API.LLVMAtomicOrdering + +is_ordered(order::AtomicOrdering) = + order ∉ (LLVM.API.LLVMAtomicOrderingNotAtomic, LLVM.API.LLVMAtomicOrderingUnordered, + LLVM.API.LLVMAtomicOrderingMonotonic) +is_acquire(order::AtomicOrdering) = + order in (LLVM.API.LLVMAtomicOrderingAcquire, LLVM.API.LLVMAtomicOrderingAcquireRelease, + LLVM.API.LLVMAtomicOrderingSequentiallyConsistent) +is_release(order::AtomicOrdering) = + order in (LLVM.API.LLVMAtomicOrderingRelease, LLVM.API.LLVMAtomicOrderingAcquireRelease, + LLVM.API.LLVMAtomicOrderingSequentiallyConsistent) + +is_atomic_memop(inst::LLVM.Instruction) = + ((inst isa LLVM.LoadInst || inst isa LLVM.StoreInst) && is_atomic(inst)) || + inst isa LLVM.AtomicRMWInst || inst isa LLVM.AtomicCmpXchgInst + +atomic_pointer(inst::LLVM.StoreInst) = operands(inst)[2] +atomic_pointer(inst::LLVM.Instruction) = operands(inst)[1] + +atomic_value_type(inst::LLVM.LoadInst) = value_type(inst) +atomic_value_type(inst::LLVM.StoreInst) = value_type(operands(inst)[1]) +atomic_value_type(inst::LLVM.Instruction) = value_type(operands(inst)[2]) + +is_volatile(inst::LLVM.Instruction) = LLVM.API.LLVMGetVolatile(inst) != 0 + +# LLVM.jl cannot name a synchronization scope, so read it from the textual form +function syncscope_name(inst::LLVM.Instruction) + m = match(r"syncscope\(\"((?:[^\"\\]|\\.)*)\"\)", string(inst)) + return m === nothing ? "system" : repr(m.captures[1]) +end + +# the ordering that covers both of a compare-exchange's orderings +# (`AtomicCmpXchgInst::getMergedOrdering`) +function merged_ordering(inst::LLVM.AtomicCmpXchgInst) + success, failure = success_ordering(inst), failure_ordering(inst) + failure == LLVM.API.LLVMAtomicOrderingSequentiallyConsistent && return failure + if failure == LLVM.API.LLVMAtomicOrderingAcquire + success == LLVM.API.LLVMAtomicOrderingMonotonic && return failure + success == LLVM.API.LLVMAtomicOrderingRelease && + return LLVM.API.LLVMAtomicOrderingAcquireRelease + end + return success +end +atomic_ordering(inst::LLVM.AtomicCmpXchgInst) = merged_ordering(inst) +atomic_ordering(inst::LLVM.Instruction) = ordering(inst) + +# the failure ordering of a compare-exchange implementing an operation with the given +# ordering (`AtomicCmpXchgInst::getStrongestFailureOrdering`) +failure_ordering_for(order::AtomicOrdering) = + order == LLVM.API.LLVMAtomicOrderingAcquireRelease ? LLVM.API.LLVMAtomicOrderingAcquire : + order == LLVM.API.LLVMAtomicOrderingRelease ? LLVM.API.LLVMAtomicOrderingMonotonic : + order + +# The operation of an `atomicrmw`. Parse it from the textual form, as the C API only knows the +# operations of the LLVM version its headers came from (LLVM 18 aborts on `uinc_wrap`). +function atomicrmw_op(inst::LLVM.AtomicRMWInst) + m = match(r"\batomicrmw\s+(?:volatile\s+)?([a-z_]+)\s", string(inst)) + m === nothing && error("Unexpected atomicrmw instruction: $inst") + return Symbol(m.captures[1]) +end + +# the value of a read-modify-write operation (`llvm::buildAtomicRMWValue`) +function atomicrmw_value!(builder::IRBuilder, op::Symbol, old::LLVM.Value, val::LLVM.Value) + T = value_type(old) + minmax(pred) = select!(builder, icmp!(builder, pred, old, val), old, val) + function intrinsic(name) + mod = LLVM.parent(LLVM.parent(position(builder))) + intr = LLVM.Intrinsic(name) + call!(builder, LLVM.FunctionType(intr, [T]), LLVM.Function(mod, intr, [T]), [old, val]) + end + op == :xchg && return val + op == :add && return add!(builder, old, val) + op == :sub && return sub!(builder, old, val) + op == :and && return and!(builder, old, val) + op == :nand && return not!(builder, and!(builder, old, val)) + op == :or && return or!(builder, old, val) + op == :xor && return xor!(builder, old, val) + op == :max && return minmax(LLVM.API.LLVMIntSGT) + op == :min && return minmax(LLVM.API.LLVMIntSLE) + op == :umax && return minmax(LLVM.API.LLVMIntUGT) + op == :umin && return minmax(LLVM.API.LLVMIntULE) + op == :fadd && return fadd!(builder, old, val) + op == :fsub && return fsub!(builder, old, val) + op == :fmax && return intrinsic("llvm.maxnum") + op == :fmin && return intrinsic("llvm.minnum") + op == :fmaximum && return intrinsic("llvm.maximum") + op == :fminimum && return intrinsic("llvm.minimum") + zero, one = ConstantInt(T, 0), ConstantInt(T, 1) + if op == :uinc_wrap + return select!(builder, icmp!(builder, LLVM.API.LLVMIntUGE, old, val), zero, + add!(builder, old, one)) + elseif op == :udec_wrap + wrap = or!(builder, icmp!(builder, LLVM.API.LLVMIntEQ, old, zero), + icmp!(builder, LLVM.API.LLVMIntUGT, old, val)) + return select!(builder, wrap, val, sub!(builder, old, one)) + elseif op == :usub_cond + return select!(builder, icmp!(builder, LLVM.API.LLVMIntUGE, old, val), + sub!(builder, old, val), old) + elseif op == :usub_sat + return select!(builder, icmp!(builder, LLVM.API.LLVMIntUGE, old, val), + sub!(builder, old, val), zero) + end + error("Unsupported atomicrmw operation: $op") +end diff --git a/src/metal.jl b/src/metal.jl index 3180b02f..5c9515c6 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -465,6 +465,19 @@ function validate_ir(job::CompilerJob{MetalCompilerTarget}, mod::LLVM.Module) # Metal never supports 128-bit integers append!(errors, check_ir_values(mod, LLVM.IntType(128))) + # atomics that `lower_atomics!` cannot lower + for f in functions(mod), bb in blocks(f), inst in instructions(bb) + reason = if is_atomic_memop(inst) + action = metal_atomic_action(job, inst) + action isa String ? action : nothing + elseif inst isa LLVM.FenceInst && metal_thread_scope(inst) === nothing + "fence with synchronization scope $(syncscope_name(inst))" + else + nothing + end + reason === nothing || push!(errors, (reason, backtrace(inst), string(inst))) + end + errors end @@ -663,6 +676,10 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L # Metal.malloc uses. rewrite_generic_null_selects!(mod) + # lower LLVM atomics to AIR atomic intrinsics (including the fences that ordered atomics + # get bracketed with on targets without ordered atomics, so this goes first) + changed = lower_atomics!(job, mod) + # the macOS 27 back-end rejects bare LLVM fences (Metal.jl#968) lower_fences!(job, mod) @@ -682,13 +699,12 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L lower_unreachable_control_flow!(job, mod) # lower LLVM intrinsics that AIR doesn't support - changed = false for f in functions(mod) changed |= lower_llvm_intrinsics!(job, f) end if changed - # lowering may have introduced additional functions marked `alwaysinline`, - # and left dead declarations of the replaced LLVM intrinsics behind + # lowering may have introduced additional functions marked `alwaysinline` (including + # the atomic expansions), and left dead declarations of replaced LLVM intrinsics behind @dispose pb=NewPMPassBuilder() begin add!(pb, AlwaysInlinerPass()) add!(pb, NewPMFunctionPassManager()) do fpm @@ -708,6 +724,752 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L return end +# lowering of LLVM atomics +# +# Metal has no LLVM back-end, so the work that `AtomicExpand` and instruction selection do for +# LLVM atomics on other targets happens here. Front-ends (Metal.jl's atomic functions, +# UnsafeAtomics and Atomix, Julia's atomic intrinsics) emit plain LLVM atomics, with an +# ordering and a synchronization scope, and `lower_atomics!` turns them into the +# `air.atomic.*` intrinsics MSL uses: +# +# - atomics on the thread's own memory become plain accesses (`demote_private_atomic!`); +# - on targets without ordered atomics (MSL < 4.1), an ordered operation becomes a relaxed one +# bracketed by fences, as `AtomicExpand` does for targets that `shouldInsertFencesForAtomic`; +# - operations that AIR cannot express are rewritten in terms of ones it can +# (see `metal_atomic_action`): floating-point loads, stores and exchanges are cast to integers, +# 8- and 16-bit operations become masked operations on the containing 32-bit word, and +# read-modify-write operations without an AIR equivalent become compare-exchange loops; +# - the remaining operations are selected to `air.atomic.*` calls (`select_atomic!`), in the +# form the target's AIR and MSL versions use. +# +# Front-ends can also call `air.atomic.*` intrinsics directly, e.g., to pass memory flags that +# LLVM cannot express. They do so in the MSL 4.1 (AIR 2.9) form, which `legalize_atomic_abi!` +# rewrites for older targets. `validate_ir` rejects the atomics that cannot be lowered (see +# `metal_atomic_action`), so the lowering can assume every atomic it sees is supported. + +# MSL memory_order values: relaxed=0, acquire=2, release=3, acq_rel=4, seq_cst=5 +metal_memory_order(order::AtomicOrdering) = + order == LLVM.API.LLVMAtomicOrderingAcquire ? 2 : + order == LLVM.API.LLVMAtomicOrderingRelease ? 3 : + order == LLVM.API.LLVMAtomicOrderingAcquireRelease ? 4 : + order == LLVM.API.LLVMAtomicOrderingSequentiallyConsistent ? 5 : 0 + +# MSL mem_flags naming the memory an ordered operation orders. LLVM orders all memory, so +# cover device and threadgroup memory, the writable address spaces LLVM code can access. +const METAL_MEM_FLAGS = 1 | 2 # mem_device | mem_threadgroup + +# The MSL thread_scope for the synchronization scope of an atomic operation or fence: thread=0, +# simdgroup=4, threadgroup=1, device=2. Scopes are spelled like the LLVM SPIR-V back-end +# does: `singlethread`, `subgroup`, `workgroup`, `device`, and the system scope (LLVM's +# default). Metal code can only synchronize with other threads on the same device (MSL has no +# scope that includes the host or other devices), so the system scope is the device scope. +# Threadgroup memory is only shared within a threadgroup, and MSL never uses a wider scope +# for it. Returns `nothing` for other scopes, which `validate_ir` rejects (like the NVPTX and +# AMDGPU back-ends do) rather than guessing what they mean. +function metal_thread_scope(inst::LLVM.Instruction, as::Union{Nothing,Int}=nothing) + ss = syncscope(inst) + scope = if ss == SyncScope("singlethread") + 0 + elseif ss == SyncScope("subgroup") + 4 + elseif ss == SyncScope("workgroup") + 1 + elseif ss == SyncScope("device") || ss == SyncScope("system") + 2 + else + return nothing + end + return as == 3 && scope == 2 ? 1 : scope +end + +# read-modify-write operations AIR has 32-bit intrinsics for, and the intrinsic names +const AIR_ATOMICRMW_OPS = Dict( + :xchg => "xchg", :add => "add.s", :sub => "sub.s", :and => "and.s", :or => "or.s", + :xor => "xor.s", :max => "max.s", :min => "min.s", :umax => "max.u", :umin => "min.u", + :fadd => "add", :fsub => "sub") + +# read-modify-write operations we can expand to compare-exchange loops +const EXPANDABLE_ATOMICRMW_OPS = (:nand, :fmax, :fmin, :fmaximum, :fminimum, :uinc_wrap, + :udec_wrap, :usub_cond, :usub_sat) + +function atomic_bits(T::LLVMType) + T isa LLVM.IntegerType && return Int(width(T)) + T isa LLVM.LLVMHalf && return 16 + T isa LLVM.LLVMBFloat && return 16 + T isa LLVM.LLVMFloat && return 32 + T isa LLVM.LLVMDouble && return 64 + T isa LLVM.PointerType && return 64 + return nothing +end + +# Does `ptr` point to the thread's own stack, i.e., is every object it can be derived from an +# `alloca`? That is the case for atomics on objects that Julia's `AllocOpt` moved to the stack, +# e.g., a non-escaping mutable struct with `@atomic` fields (GPUCompiler.jl#934). Metal cannot +# express those (MSL only has atomics on device and threadgroup memory), but they don't need +# to be atomic: no other thread can access a thread's stack, even when it has a pointer to it +# (thread memory is private to every thread), so plain accesses behave the same. Anything this +# cannot trace back to an `alloca`, e.g., a function argument or a loaded pointer, is not +# known to be private. +function is_thread_private(ptr::LLVM.Value) + seen = Set{LLVM.Value}() + worklist = LLVM.Value[ptr] + while !isempty(worklist) + val = pop!(worklist) + val in seen && continue + push!(seen, val) + if val isa LLVM.AllocaInst + continue + elseif val isa LLVM.GetElementPtrInst || val isa LLVM.BitCastInst + push!(worklist, operands(val)[1]) + elseif val isa LLVM.PHIInst + append!(worklist, first.(LLVM.incoming(val))) + elseif val isa LLVM.SelectInst + push!(worklist, operands(val)[2], operands(val)[3]) + else + return false + end + end + return true +end + +# How to lower `inst`, an atomic memory operation, for the job's target. Like the rule tables +# of LLVM's legalizers, the rules are tried in order and the first that applies decides. Returns +# the action, or the reason why the operation cannot be lowered (which `validate_ir` reports): +# +# - `:demote`: an atomic on the thread's own memory, which becomes plain accesses; +# - `:cast`: a floating-point load, store or exchange, which becomes an integer one (like the +# default `TargetLowering::shouldCast*InIR`), so that AIR's integer intrinsics can be used; +# - `:partword`: an 8- or 16-bit operation, which becomes a masked operation on the containing +# 32-bit word (`AtomicExpand`'s `expandPartwordAtomicRMW` and `expandPartwordCmpXchg`); +# - `:cmpxchg_loop`: a read-modify-write operation without an AIR equivalent, which becomes a +# compare-exchange loop (`AtomicExpand`'s `insertRMWCmpXchgLoop`); +# - `:select`: an operation AIR can express, which becomes an `air.atomic.*` call. +function metal_atomic_action(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + inst::LLVM.Instruction) + target = job.config.target + op = inst isa LLVM.AtomicRMWInst ? atomicrmw_op(inst) : nothing + if op !== nothing && !haskey(AIR_ATOMICRMW_OPS, op) && !(op in EXPANDABLE_ATOMICRMW_OPS) + return "atomicrmw $op operation" + end + + is_thread_private(atomic_pointer(inst)) && return :demote + + as = addrspace(value_type(atomic_pointer(inst))) + if as != 1 && as != 3 + return "atomic operation in address space $as (Metal only supports atomics on device and threadgroup memory)" + end + + T = atomic_value_type(inst) + bits = atomic_bits(T) + if bits === nothing || !(bits in (8, 16, 32, 64)) + return "atomic operation on a $(string(T)) value" + end + + # AIR's only 64-bit atomics are umin/umax on device memory, which don't return the old + # value (and need an Apple8 GPU). There is no 64-bit compare-exchange to emulate others. + if bits == 64 && !(op in (:umin, :umax) && T isa LLVM.IntegerType && as == 1 && + isempty(uses(inst))) + return "64-bit atomic operation (Metal only supports atomic 64-bit umin and umax on device memory, without using the result)" + end + + if alignment(inst) < bits ÷ 8 + return "misaligned atomic operation" + end + + if metal_thread_scope(inst) === nothing + return "atomic operation with synchronization scope $(syncscope_name(inst))" + end + + # without ordered atomics, orderings are implemented with fences (from MSL 3.2) + if is_ordered(atomic_ordering(inst)) && target.metal < v"3.2" + return "ordered atomic operation (Metal $(target.metal) only supports relaxed atomics)" + end + + if T isa LLVM.FloatingPointType && (op === nothing || op == :xchg) && + !(inst isa LLVM.AtomicCmpXchgInst) + return :cast + end + bits < 32 && return :partword + # (threadgroup floating-point add and subtract need MSL 4.1) + if op in EXPANDABLE_ATOMICRMW_OPS || + (op in (:fadd, :fsub) && as == 3 && target.metal < v"4.1") + return :cmpxchg_loop + end + return :select +end + +# Build an expansion that needs control flow as an internal, always-inlined function, and +# replace `inst` with a call to it: LLVM.jl cannot split basic blocks. `body(builder, f, +# params...)` emits the function's code and returns its result, and the inliner that runs at +# the end of `lower_air!` puts the code in place. +function outline_atomic!(body, mod::LLVM.Module, inst::LLVM.Instruction, + args::Vector{<:LLVM.Value}) + T_ret = value_type(inst) + ft = LLVM.FunctionType(T_ret, map(value_type, args)) + f = LLVM.Function(mod, "julia.air.atomic_expansion", ft) + linkage!(f, LLVM.API.LLVMInternalLinkage) + push!(function_attributes(f), EnumAttribute("alwaysinline")) + @dispose builder=IRBuilder() begin + position!(builder, BasicBlock(f, "entry")) + result = body(builder, f, parameters(f)...) + T_ret == LLVM.VoidType() ? ret!(builder) : ret!(builder, result) + + position!(builder, inst) + debuglocation!(builder, inst) + call = call!(builder, ft, f, args) + T_ret == LLVM.VoidType() || replace_uses!(inst, call) + end + erase!(inst) + return +end + +function set_atomic!(inst::LLVM.Instruction, order::AtomicOrdering, scope::SyncScope, + volatile::Bool=false) + ordering!(inst, order) + syncscope!(inst, scope) + volatile && LLVM.API.LLVMSetVolatile(inst, true) + return inst +end + +# Replace an atomic operation on the thread's own memory (see `is_thread_private`) by plain +# accesses. Its ordering and scope don't matter either: no other thread can observe the memory +# it accesses, so it cannot synchronize with any. +function demote_private_atomic!(inst::LLVM.Instruction) + ptr = atomic_pointer(inst) + T = atomic_value_type(inst) + volatile = is_volatile(inst) + @dispose builder=IRBuilder() begin + position!(builder, inst) + debuglocation!(builder, inst) + function plain_load() + ld = load!(builder, T, ptr) + alignment!(ld, alignment(inst)) + volatile && LLVM.API.LLVMSetVolatile(ld, true) + ld + end + function plain_store(val) + st = store!(builder, val, ptr) + alignment!(st, alignment(inst)) + volatile && LLVM.API.LLVMSetVolatile(st, true) + st + end + if inst isa LLVM.LoadInst + replace_uses!(inst, plain_load()) + elseif inst isa LLVM.StoreInst + plain_store(operands(inst)[1]) + elseif inst isa LLVM.AtomicRMWInst + old = plain_load() + plain_store(atomicrmw_value!(builder, atomicrmw_op(inst), old, operands(inst)[2])) + replace_uses!(inst, old) + else + # compare-exchange: store the new value if the old one matches, else the old one + cmp, new = operands(inst)[2:3] + old = plain_load() + success = icmp!(builder, LLVM.API.LLVMIntEQ, old, cmp) + plain_store(select!(builder, success, new, old)) + result = insert_value!(builder, UndefValue(value_type(inst)), old, 0) + replace_uses!(inst, insert_value!(builder, result, success, 1)) + end + end + erase!(inst) + return +end + +# Emit a loop that atomically replaces the 32-bit word at `ptr` by `update(builder, word)` +# using compare-exchange, returning the word the successful exchange replaced +# (`AtomicExpand`'s `insertRMWCmpXchgLoop`). +function emit_cmpxchg_loop!(update, builder::IRBuilder, f::LLVM.Function, ptr::LLVM.Value, + order::AtomicOrdering, scope::SyncScope, volatile::Bool) + T_word = LLVM.Int32Type() + entry = position(builder) + init = load!(builder, T_word, ptr) + alignment!(init, 4) + set_atomic!(init, LLVM.API.LLVMAtomicOrderingMonotonic, scope, volatile) + loop = BasicBlock(f, "atomicrmw.start") + done = BasicBlock(f, "atomicrmw.end") + br!(builder, loop) + + position!(builder, loop) + loaded = phi!(builder, T_word, "loaded") + pair = atomic_cmpxchg!(builder, ptr, loaded, update(builder, loaded), order, + failure_ordering_for(order), scope) + volatile && LLVM.API.LLVMSetVolatile(pair, true) + word = extract_value!(builder, pair, 0) + br!(builder, extract_value!(builder, pair, 1), done, loop) + push!(LLVM.incoming(loaded), (init, entry)) + push!(LLVM.incoming(loaded), (word, loop)) + + position!(builder, done) + return word +end + +# The 32-bit word containing an 8- or 16-bit value at `ptr`, and the position of the value +# in it (`AtomicExpand`'s `createMaskInstrs`; Metal is little-endian). Like `AtomicExpand`, +# this assumes that whole word can be accessed: true for device buffers, which Metal allocates +# in pages, and for threadgroup arrays, which Metal.jl aligns and pads to 4 bytes. +function partword_layout!(builder::IRBuilder, ptr::LLVM.Value, bits::Int) + as = addrspace(value_type(ptr)) + T_i8, T_i32, T_i64 = LLVM.Int8Type(), LLVM.Int32Type(), LLVM.Int64Type() + bytes = bitcast!(builder, ptr, LLVM.PointerType(T_i8, as)) + offset = and!(builder, ptrtoint!(builder, bytes, T_i64), ConstantInt(T_i64, 3)) + word = gep!(builder, T_i8, bytes, [neg!(builder, offset)]) + word = bitcast!(builder, word, LLVM.PointerType(T_i32, as)) + shift = shl!(builder, trunc!(builder, offset, T_i32), ConstantInt(T_i32, 3)) + mask = shl!(builder, ConstantInt(T_i32, (1 << bits) - 1), shift) + return (; word, shift, mask, inv_mask=not!(builder, mask)) +end +partword_extract!(builder, layout, word, T) = + trunc!(builder, lshr!(builder, word, layout.shift), T) +partword_insert!(builder, layout, word, val) = + or!(builder, and!(builder, word, layout.inv_mask), + shl!(builder, zext!(builder, val, LLVM.Int32Type()), layout.shift)) + +# The operands of an atomic operation that its expansions need, with the ordering it is lowered +# with (see `lowered_ordering`). +function atomic_operands(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + inst::LLVM.Instruction) + T = atomic_value_type(inst) + return (; ptr=atomic_pointer(inst), T, bits=atomic_bits(T), + op=inst isa LLVM.AtomicRMWInst ? atomicrmw_op(inst) : nothing, + order=lowered_ordering(job, atomic_ordering(inst)), scope=syncscope(inst), + volatile=is_volatile(inst)) +end + +# Cast a floating-point load, store or exchange to an integer one, returning the new operation. +function cast_atomic_to_int!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + inst::LLVM.Instruction) + (; ptr, T, bits, order, scope, volatile) = atomic_operands(job, inst) + T_int = LLVM.IntType(bits) + @dispose builder=IRBuilder() begin + position!(builder, inst) + debuglocation!(builder, inst) + int_ptr = bitcast!(builder, ptr, LLVM.PointerType(T_int, addrspace(value_type(ptr)))) + new = if inst isa LLVM.LoadInst + ld = load!(builder, T_int, int_ptr) + alignment!(ld, alignment(inst)) + set_atomic!(ld, order, scope, volatile) + replace_uses!(inst, bitcast!(builder, ld, T)) + ld + elseif inst isa LLVM.StoreInst + st = store!(builder, bitcast!(builder, operands(inst)[1], T_int), int_ptr) + alignment!(st, alignment(inst)) + set_atomic!(st, order, scope, volatile) + else + rmw = atomic_rmw!(builder, LLVM.API.LLVMAtomicRMWBinOpXchg, int_ptr, + bitcast!(builder, operands(inst)[2], T_int), order, scope) + alignment!(rmw, alignment(inst)) + volatile && LLVM.API.LLVMSetVolatile(rmw, true) + replace_uses!(inst, bitcast!(builder, rmw, T)) + rmw + end + erase!(inst) + return new + end +end + +# Expand a 32-bit read-modify-write operation to a compare-exchange loop on the word. +function expand_to_cmpxchg_loop!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + mod::LLVM.Module, inst::LLVM.Instruction) + (; ptr, T, op, order, scope, volatile) = atomic_operands(job, inst) + T_i32 = LLVM.Int32Type() + outline_atomic!(mod, inst, [ptr, operands(inst)[2]]) do builder, f, ptr, val + word_ptr = bitcast!(builder, ptr, LLVM.PointerType(T_i32, addrspace(value_type(ptr)))) + word = emit_cmpxchg_loop!(builder, f, word_ptr, order, scope, volatile) do builder, loaded + old = bitcast!(builder, loaded, T) + bitcast!(builder, atomicrmw_value!(builder, op, old, val), T_i32) + end + bitcast!(builder, word, T) + end + return +end + +# 8- and 16-bit atomics as masked operations on the containing 32-bit word +# (`AtomicExpand`'s `expandPartwordAtomicRMW` and `expandPartwordCmpXchg`) +function expand_partword_atomic!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + mod::LLVM.Module, inst::LLVM.Instruction) + (; ptr, T, bits, op, order, scope, volatile) = atomic_operands(job, inst) + T_i32 = LLVM.Int32Type() + + # operations that don't need a loop: loads, and bitwise operations that leave the rest + # of the word unchanged + if inst isa LLVM.LoadInst || op in (:and, :or, :xor) + @dispose builder=IRBuilder() begin + position!(builder, inst) + debuglocation!(builder, inst) + layout = partword_layout!(builder, ptr, bits) + word = if inst isa LLVM.LoadInst + ld = load!(builder, T_i32, layout.word) + alignment!(ld, 4) + set_atomic!(ld, order, scope, volatile) + else + val = shl!(builder, zext!(builder, operands(inst)[2], T_i32), layout.shift) + op == :and && (val = or!(builder, val, layout.inv_mask)) + binop = op == :and ? LLVM.API.LLVMAtomicRMWBinOpAnd : + op == :or ? LLVM.API.LLVMAtomicRMWBinOpOr : + LLVM.API.LLVMAtomicRMWBinOpXor + rmw = atomic_rmw!(builder, binop, layout.word, val, order, scope) + alignment!(rmw, 4) + volatile && LLVM.API.LLVMSetVolatile(rmw, true) + rmw + end + replace_uses!(inst, partword_extract!(builder, layout, word, T)) + end + erase!(inst) + return + end + + if inst isa LLVM.AtomicCmpXchgInst + cmp, new = operands(inst)[2:3] + success = lowered_ordering(job, success_ordering(inst)) + failure = lowered_ordering(job, failure_ordering(inst)) + T_result = value_type(inst) + outline_atomic!(mod, inst, [ptr, cmp, new]) do builder, f, ptr, cmp, new + layout = partword_layout!(builder, ptr, bits) + new_shifted = shl!(builder, zext!(builder, new, T_i32), layout.shift) + cmp_shifted = shl!(builder, zext!(builder, cmp, T_i32), layout.shift) + init = load!(builder, T_i32, layout.word) + alignment!(init, 4) + set_atomic!(init, LLVM.API.LLVMAtomicOrderingMonotonic, scope, volatile) + init_rest = and!(builder, init, layout.inv_mask) + entry = position(builder) + loop = BasicBlock(f, "partword.cmpxchg.loop") + failed = BasicBlock(f, "partword.cmpxchg.failure") + done = BasicBlock(f, "partword.cmpxchg.end") + br!(builder, loop) + + # retry as long as the exchange only failed because of the rest of the word + position!(builder, loop) + rest = phi!(builder, T_i32, "rest") + pair = atomic_cmpxchg!(builder, layout.word, or!(builder, rest, cmp_shifted), + or!(builder, rest, new_shifted), success, failure, scope) + volatile && LLVM.API.LLVMSetVolatile(pair, true) + word = extract_value!(builder, pair, 0) + ok = extract_value!(builder, pair, 1) + br!(builder, ok, done, failed) + + position!(builder, failed) + new_rest = and!(builder, word, layout.inv_mask) + br!(builder, icmp!(builder, LLVM.API.LLVMIntNE, rest, new_rest), loop, done) + push!(LLVM.incoming(rest), (init_rest, entry)) + push!(LLVM.incoming(rest), (new_rest, failed)) + + position!(builder, done) + result = insert_value!(builder, UndefValue(T_result), + partword_extract!(builder, layout, word, T), 0) + insert_value!(builder, result, ok, 1) + end + return + end + + # everything else becomes a compare-exchange loop on the word; the value may be a + # floating-point one, for the arithmetic read-modify-write operations + val = inst isa LLVM.StoreInst ? operands(inst)[1] : operands(inst)[2] + op = something(op, :xchg) # a store is an exchange with an ignored result + # compare-exchange needs at least monotonic (like `AtomicExpand`'s `expandAtomicStoreToXChg`) + order == LLVM.API.LLVMAtomicOrderingUnordered && (order = LLVM.API.LLVMAtomicOrderingMonotonic) + is_store = inst isa LLVM.StoreInst + T_int = LLVM.IntType(bits) + outline_atomic!(mod, inst, [ptr, val]) do builder, f, ptr, val + layout = partword_layout!(builder, ptr, bits) + word = emit_cmpxchg_loop!(builder, f, layout.word, order, scope, + volatile) do builder, loaded + old = bitcast!(builder, partword_extract!(builder, layout, loaded, T_int), T) + new = bitcast!(builder, atomicrmw_value!(builder, op, old, val), T_int) + partword_insert!(builder, layout, loaded, new) + end + is_store ? nothing : + bitcast!(builder, partword_extract!(builder, layout, word, T_int), T) + end + return +end + +# On targets without ordered atomics (MSL < 4.1), bracket the operation with fences +# (`AtomicExpand`'s `bracketInstWithFences` with the default `emitLeadingFence` and +# `emitTrailingFence`), which `lower_fences!` turns into `air.atomic.fence` calls; the +# operation itself is then lowered as a relaxed one (see `lowered_ordering`). The ordering is +# not reset here, as the C API cannot do so for `atomicrmw` before LLVM 18. +function insert_atomic_fences!(inst::LLVM.Instruction) + order = atomic_ordering(inst) + is_ordered(order) || return false + scope = syncscope(inst) + @dispose builder=IRBuilder() begin + if is_release(order) && !(inst isa LLVM.LoadInst) + position!(builder, inst) + debuglocation!(builder, inst) + fence!(builder, order, scope) + end + if is_acquire(order) + position!(builder, nextinst(inst)) + debuglocation!(builder, inst) + fence!(builder, order, scope) + end + end + return true +end + +# LLVM requires that threads repeatedly loading an address monotonically eventually see the +# stores of other threads, but Apple's compiler emits relaxed loads of device memory as cached +# loads: on an M1, a spin loop that relaxed-loads a flag another threadgroup sets never sees +# the store (not in 20M iterations; within a threadgroup it does, as it shares the cache). +# Acquire loads invalidate the cache after loading, so lower device-scope monotonic loads of +# device memory as acquire ones (which on targets without ordered atomics adds a fence; before +# MSL 3.2, there are no fences to lower that to). +function strengthen_relaxed_load!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + inst::LLVM.Instruction) + job.config.target.metal >= v"3.2" && inst isa LLVM.LoadInst && + atomic_ordering(inst) == LLVM.API.LLVMAtomicOrderingMonotonic && + addrspace(value_type(atomic_pointer(inst))) == 1 && + metal_thread_scope(inst, 1) == 2 || return false + ordering!(inst, LLVM.API.LLVMAtomicOrderingAcquire) + return true +end + +# With ordered atomics (MSL ≥ 4.1), follow a sequentially-consistent store by a +# sequentially-consistent fence (AtomicExpand's `shouldInsertTrailingSeqCstFenceForAtomicStore`, +# which AArch64 uses for MSVC). Apple compiles such a store as a release store that doesn't +# wait for the write, so a later sequentially-consistent load can be performed first: +# `store x; load y` and `store y; load x` in two threads both return the old values (store +# buffering, ~0.5% of the time on an M1), which sequential consistency forbids. +# Read-modify-writes and compare-exchanges wait for their result and don't need this, and +# fences before loads instead would cost more (loads are more common than stores). +function insert_trailing_seq_cst_fence!(inst::LLVM.Instruction) + inst isa LLVM.StoreInst && + atomic_ordering(inst) == LLVM.API.LLVMAtomicOrderingSequentiallyConsistent || + return false + @dispose builder=IRBuilder() begin + position!(builder, nextinst(inst)) + debuglocation!(builder, inst) + fence!(builder, LLVM.API.LLVMAtomicOrderingSequentiallyConsistent, syncscope(inst)) + end + return true +end + +# The ordering to lower an atomic operation with: without ordered atomics (MSL < 4.1), the +# fences `insert_atomic_fences!` added provide the ordering, and the operation is relaxed. +lowered_ordering(@nospecialize(job::CompilerJob{MetalCompilerTarget}), order::AtomicOrdering) = + job.config.target.metal < v"4.1" ? LLVM.API.LLVMAtomicOrderingMonotonic : order + +function air_atomic_function(mod::LLVM.Module, name::String, ft::LLVM.FunctionType) + if haskey(functions(mod), name) + f = functions(mod)[name] + function_type(f) == ft || + error("Conflicting declarations of $name: $(function_type(f)) and $ft") + return f + end + f = LLVM.Function(mod, name, ft) + # as Apple declares them (not `argmemonly` or `readonly`: they order other memory) + for attr in ("mustprogress", "nounwind", "willreturn") + push!(function_attributes(f), EnumAttribute(attr, 0)) + end + return f +end + +# Replace an atomic operation by the equivalent `air.atomic.*` call, in the form the +# target's AIR and MSL versions use. +function select_atomic!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + mod::LLVM.Module, inst::LLVM.Instruction) + target = job.config.target + ptr = atomic_pointer(inst) + T_ptr = value_type(ptr) + as = addrspace(T_ptr) + T = atomic_value_type(inst) + mem = as == 1 ? "global" : "local" + suffix = T isa LLVM.LLVMFloat ? "f32" : "i$(width(T))" + T_i32, T_i1 = LLVM.Int32Type(), LLVM.Int1Type() + + # the operands after the memory order(s): scope, flags (from AIR 2.9), volatile + order = lowered_ordering(job, atomic_ordering(inst)) + # (the scope of a relaxed load of device memory only matters to whether we strengthen it, + # so use the device scope, like MSL does: it doesn't change the code, but it is recorded) + relaxed_device_load = inst isa LLVM.LoadInst && as == 1 && + atomic_ordering(inst) == LLVM.API.LLVMAtomicOrderingMonotonic + scope = ConstantInt(T_i32, relaxed_device_load ? 2 : metal_thread_scope(inst, as)) + flags = ConstantInt(T_i32, is_ordered(order) ? METAL_MEM_FLAGS : 0) + # MSL sets the volatile bit on every atomic before 4.1, but since then only on atomics + # of `volatile` objects. Without it, the back-end treats a load like a plain one (e.g., + # a relaxed load of memory the kernel doesn't write is hoisted into the uniform preamble, + # out of any spin loop), and a read-modify-write that doesn't change memory (e.g., adding + # 0) becomes such a load. LLVM atomics allow neither, so always set it on loads and + # read-modify-writes; stores and compare-exchanges are compiled the same either way. + volatile = ConstantInt(T_i1, target.metal < v"4.1" || is_volatile(inst) || + inst isa LLVM.LoadInst || inst isa LLVM.AtomicRMWInst) + trailing_types = target.air >= v"2.9" ? [T_i32, T_i32, T_i1] : [T_i32, T_i1] + trailing = target.air >= v"2.9" ? [scope, flags, volatile] : [scope, volatile] + memory_order(order) = ConstantInt(T_i32, metal_memory_order(lowered_ordering(job, order))) + + @dispose builder=IRBuilder() begin + position!(builder, inst) + debuglocation!(builder, inst) + if inst isa LLVM.LoadInst + ft = LLVM.FunctionType(T, [T_ptr, T_i32, trailing_types...]) + f = air_atomic_function(mod, "air.atomic.$mem.load.$suffix", ft) + new = call!(builder, ft, f, [ptr, memory_order(order), trailing...]) + elseif inst isa LLVM.StoreInst + ft = LLVM.FunctionType(LLVM.VoidType(), [T_ptr, T, T_i32, trailing_types...]) + f = air_atomic_function(mod, "air.atomic.$mem.store.$suffix", ft) + new = call!(builder, ft, f, [ptr, operands(inst)[1], memory_order(order), + trailing...]) + elseif inst isa LLVM.AtomicRMWInst + op = AIR_ATOMICRMW_OPS[atomicrmw_op(inst)] + # AIR's 64-bit min/max don't return the old value + T_ret = suffix == "i64" ? LLVM.VoidType() : T + ft = LLVM.FunctionType(T_ret, [T_ptr, T, T_i32, trailing_types...]) + f = air_atomic_function(mod, "air.atomic.$mem.$op.$suffix", ft) + new = call!(builder, ft, f, [ptr, operands(inst)[2], memory_order(order), + trailing...]) + else + # AIR only has a weak compare-exchange, which takes the expected value by + # reference. Like MSL, derive the success flag from the returned old value. + cmp, desired = operands(inst)[2:3] + fn = LLVM.parent(LLVM.parent(inst)) + expected = @dispose entry_builder=IRBuilder() begin + position!(entry_builder, first(instructions(first(blocks(fn))))) + alloca!(entry_builder, T) + end + store!(builder, cmp, expected) + ft = LLVM.FunctionType(T, [T_ptr, value_type(expected), T, T_i32, T_i32, + trailing_types...]) + f = air_atomic_function(mod, "air.atomic.$mem.cmpxchg.weak.$suffix", ft) + old = call!(builder, ft, f, + [ptr, expected, desired, memory_order(success_ordering(inst)), + memory_order(failure_ordering(inst)), trailing...]) + success = icmp!(builder, LLVM.API.LLVMIntEQ, old, cmp) + new = insert_value!(builder, UndefValue(value_type(inst)), old, 0) + new = insert_value!(builder, new, success, 1) + end + isempty(uses(inst)) || replace_uses!(inst, new) + end + erase!(inst) + return +end + +# the number of operands of an `air.atomic.*` intrinsic in its MSL 4.1 (AIR 2.9) form +function air_atomic_arity(name::String) + occursin(".cmpxchg.", name) && return 8 + occursin(".load.", name) && return 5 + return 6 +end + +# Rewrite the `air.atomic.*` calls front-ends emit (in the MSL 4.1 form) into the form the +# target uses: AIR 2.9 introduced the memory flags operand, and before MSL 4.1 the volatile bit +# is always set, only relaxed orderings are supported and the flags must be zero. Also tell the +# downgrader the element types of the pointer operands, which LLVM cannot infer. +function legalize_atomic_abi!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + mod::LLVM.Module) + target = job.config.target + changed = false + for f in collect(functions(mod)) + fn = LLVM.name(f) + m = match(r"^air\.atomic\.(global|local)\..+\.(i32|f32|i64)$", fn) + m === nothing && continue + + # pointer element types for the typed-pointer downgrader + T = m.captures[2] == "f32" ? LLVM.FloatType() : LLVM.IntType(parse(Int, m.captures[2][2:end])) + mds = [] + for (i, param) in enumerate(parameters(function_type(f))) + i <= 2 && param isa LLVM.PointerType || continue + push!(mds, ConstantInt(Int32(i - 1))) + push!(mds, null(T)) + end + metadata(f)["arg_eltypes"] = MDNode(mds) + + target.metal >= v"4.1" && continue + nparams = length(parameters(function_type(f))) + # older front-ends rewrote their calls for the target themselves + nparams == air_atomic_arity(fn) || continue + + is_cmpxchg = occursin(".cmpxchg.", fn) + calls = [user(u)::LLVM.CallInst for u in uses(f)] + for call in calls + args = collect(arguments(call)) + orders = is_cmpxchg ? args[end-4:end-3] : args[end-3:end-3] + if !all(o -> o isa ConstantInt && convert(Int, o) == 0, [orders..., args[end-1]]) + error("$fn with an ordering or memory flags requires MSL 4.1; " * + "the target is MSL $(target.metal)") + end + end + changed |= !isempty(calls) + + if target.air >= v"2.9" + # the flags operand exists, but the volatile bit must be set + for call in calls + # LLVM models the callee as the final operand, after all arguments + operands(call)[end-1] = ConstantInt(true) + end + continue + end + + # before AIR 2.9, there is no flags operand + ft = function_type(f) + params = parameters(ft) + legacy_ft = LLVM.FunctionType(LLVM.return_type(ft), [params[1:end-2]..., params[end]]) + LLVM.name!(f, fn * ".msl41") + legacy_f = LLVM.Function(mod, fn, legacy_ft) + for attr in collect(function_attributes(f)) + push!(function_attributes(legacy_f), attr) + end + metadata(legacy_f)["arg_eltypes"] = metadata(f)["arg_eltypes"] + for call in calls + args = collect(arguments(call)) + @dispose builder=IRBuilder() begin + position!(builder, call) + debuglocation!(builder, call) + new = call!(builder, legacy_ft, legacy_f, + [args[1:end-2]..., ConstantInt(true)]) + replace_uses!(call, new) + end + erase!(call) + end + erase!(f) + end + return changed +end + +function lower_atomics!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), + mod::LLVM.Module) + # first rewrite the intrinsic calls front-ends emitted, so that the ones we select below + # agree with them + changed = legalize_atomic_abi!(job, mod) + + atomics = [inst for f in functions(mod) for bb in blocks(f) for inst in instructions(bb) + if is_atomic_memop(inst)] + isempty(atomics) && return changed + + for inst in atomics + action = metal_atomic_action(job, inst) + # (`validate_ir` rejects unsupported operations; without validation, select them as-is) + action isa String && continue + if action === :demote + demote_private_atomic!(inst) + continue + end + strengthen_relaxed_load!(job, inst) + if job.config.target.metal < v"4.1" + insert_atomic_fences!(inst) + else + insert_trailing_seq_cst_fence!(inst) + end + if action === :cast + inst = cast_atomic_to_int!(job, inst) + action = metal_atomic_action(job, inst) # e.g. a half-precision load is partword + end + if action === :partword + expand_partword_atomic!(job, mod, inst) + elseif action === :cmpxchg_loop + expand_to_cmpxchg_loop!(job, mod, inst) + end + end + + # select the atomics that are left, including the ones the expansions introduced + for f in functions(mod), bb in blocks(f), inst in collect(instructions(bb)) + is_atomic_memop(inst) && select_atomic!(job, mod, inst) + end + + # attach element type metadata to the declarations we introduced + legalize_atomic_abi!(job, mod) + return true +end + # Before LLVM 18, `ordering(inst)` calls `LLVMGetOrdering`, which incorrectly casts fences # to AtomicRMWInst. Use the stable textual form on all versions to keep this workaround tested. function fence_ordering(inst::LLVM.FenceInst) @@ -724,14 +1486,12 @@ end # Lower LLVM fences to air.atomic.fence(flags, order, scope), as MSL's atomic_thread_fence # does. Bare fences from Julia's atomic_fence crash the macOS 27 AGX back-end (Metal.jl#968). # -# MSL memory_order values are acquire=2, release=3, acq_rel=4, seq_cst=5. Metal 3.2-4.0 -# only supports relaxed/seq_cst fences, so strengthen other orderings to seq_cst. Before -# Metal 3.2 the intrinsic is unavailable; retain the bare fence. Such targets still -# require a back-end that accepts bare fences. +# Metal 3.2-4.0 only supports relaxed/seq_cst fences, so strengthen other orderings to +# seq_cst. Before Metal 3.2 the intrinsic is unavailable; retain the bare fence. Such targets +# still require a back-end that accepts bare fences. # -# Cover device and threadgroup memory (mem_flags=1|2), the shared writable LLVM address spaces. -# Use thread scope (0) for singlethread and device scope (2) otherwise. Metal has no -# system-wide scope: this only synchronizes threads on the same device. +# Cover device and threadgroup memory (`METAL_MEM_FLAGS`), the shared writable LLVM address +# spaces, and map the scope like for atomics (`metal_thread_scope`). function lower_fences!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::LLVM.Module) metal = job.config.target.metal metal >= v"3.2" || return false @@ -751,28 +1511,15 @@ function lower_fences!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod end for inst in worklist - order = if metal < v"4.1" - 5 # seq_cst - else - ord = fence_ordering(inst) - if ord == LLVM.API.LLVMAtomicOrderingAcquire - 2 - elseif ord == LLVM.API.LLVMAtomicOrderingRelease - 3 - elseif ord == LLVM.API.LLVMAtomicOrderingAcquireRelease - 4 - else - 5 # seq_cst (the only other ordering LLVM allows on a fence) - end - end - scope = syncscope(inst) == SyncScope("singlethread") ? 0 : 2 - flags = 1 | 2 # device | threadgroup + order = metal < v"4.1" ? 5 : metal_memory_order(fence_ordering(inst)) + # (`validate_ir` rejects unknown scopes; without validation, use the widest one) + scope = something(metal_thread_scope(inst), 2) @dispose builder=IRBuilder() begin position!(builder, inst) debuglocation!(builder, inst) call!(builder, fence_ft, fence_fn, - [ConstantInt(T_int32, flags), ConstantInt(T_int32, order), + [ConstantInt(T_int32, METAL_MEM_FLAGS), ConstantInt(T_int32, order), ConstantInt(T_int32, scope)]) end erase!(inst) @@ -2542,19 +3289,11 @@ function annotate_air_intrinsics!(@nospecialize(job::CompilerJob), mod::LLVM.Mod elseif match(r"^air.(fast_)?sincos", fn) !== nothing add_param_attributes(2, "nocapture", "writeonly") - # atomics - elseif match(r"air.atomic.(local|global).load", fn) !== nothing - # TODO: "memory(argmem: read)" on LLVM 16+ - add_fn_attributes("argmemonly", "readonly", "nounwind") - elseif match(r"air.atomic.(local|global).store", fn) !== nothing - # TODO: "memory(argmem: write)" on LLVM 16+ - add_fn_attributes("argmemonly", "writeonly", "nounwind") - elseif match(r"air.atomic.(local|global).(xchg|cmpxchg)", fn) !== nothing - # TODO: "memory(argmem: readwrite)" on LLVM 16+ - add_fn_attributes("argmemonly", "nounwind") - elseif match(r"^air.atomic.(local|global).(add|sub|min|max|and|or|xor)", fn) !== nothing - # TODO: "memory(argmem: readwrite)" on LLVM 16+ - add_fn_attributes("argmemonly", "nounwind") + # atomics: as Apple declares them. not `argmemonly` or `readonly`, which would let + # LLVM hoist an atomic load out of a spin loop, or move other memory accesses across + # an ordered atomic. + elseif match(r"^air.atomic.(local|global)\.", fn) !== nothing + add_fn_attributes("mustprogress", "nounwind", "willreturn") # simdgroup elseif match(r"air.simdgroup_matrix_8x8_init_filled", fn) !== nothing diff --git a/test/metal.jl b/test/metal.jl index d5d5af90..baf78780 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1902,9 +1902,9 @@ end end @testset "LLVM atomics" begin - # atomics in user code (e.g. UnsafeAtomics' `load`/`store!`) must reach the back-end, - # while Julia's `unordered` heap-reference accesses, which AIR cannot express when they - # are of pointers or outside device and threadgroup memory, become plain ones + # atomics in user code (e.g. UnsafeAtomics' `load`/`store!`) keep their ordering, while + # Julia's `unordered` heap-reference accesses, which AIR cannot express when they are of + # pointers or outside device and threadgroup memory, become plain ones function kernel(p::Core.LLVMPtr{Int32,1}, q::Core.LLVMPtr{Int32,1}) r = reinterpret(Ptr{Ptr{Int32}}, p) y = Core.Intrinsics.atomic_pointerref(r, :unordered) @@ -1913,6 +1913,10 @@ end Core.Intrinsics.atomic_pointerset(reinterpret(Ptr{Int32}, q), x, :release) return end + source = methodinstance(typeof(kernel), Tuple{Core.LLVMPtr{Int32,1}, Core.LLVMPtr{Int32,1}}, + Base.get_world_counter()) + target = MetalCompilerTarget(; macos=v"27", metal=v"4.1", air=v"2.9") + job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true)) @test @filecheck begin @check_not "unordered" @@ -1921,15 +1925,534 @@ end @check "store atomic i32" @check_same "release" @check_not "unordered" - Metal.code_native(kernel, Tuple{Core.LLVMPtr{Int32,1}, Core.LLVMPtr{Int32,1}}; - kernel=true, dump_module=true) + GPUCompiler.code_llvm(job; dump_module=true) + end + @test @filecheck begin + @check "call i32 @air.atomic.global.load.i32({{.+}}, i32 2, i32 2, i32 3, i1 true)" + @check "call void @air.atomic.global.store.i32({{.+}}, i32 3, i32 2, i32 3, i1 false)" + GPUCompiler.code_native(job; dump_module=true) + end +end + +# lower the atomics (and fences) of textual IR for a Metal target, like `lower_air!` +function lower_metal_atomics(ir::String; metal, air) + source = methodinstance(typeof(identity), Tuple{Int}, Base.get_world_counter()) + target = MetalCompilerTarget(; macos=v"27", metal, air) + job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true)) + Context(; opaque_pointers=true) do ctx + mod = parse(LLVM.Module, ir) + errors = GPUCompiler.validate_ir(job, mod) + isempty(errors) || return join(first.(errors), "\n") + GPUCompiler.lower_atomics!(job, mod) + GPUCompiler.lower_fences!(job, mod) + @dispose pb=NewPMPassBuilder() begin + add!(pb, AlwaysInlinerPass()) + add!(pb, StripDeadPrototypesPass()) + run!(pb, mod) + end + verify(mod) + string(mod) + end +end + +@testset "atomic lowering" begin + kernel(body; args="ptr addrspace(1) %p, ptr addrspace(3) %t") = """ + define void @f($args) { + $body + ret void + } + """ + targets = ((v"3.2", v"2.7"), (v"4.0", v"2.8"), (v"4.0", v"2.9"), (v"4.1", v"2.9")) + + @testset "orderings (Metal $metal, AIR $air)" for (metal, air) in targets + # MSL 4.1 has ordered atomics that order device and threadgroup memory (flags=3), + # but sequentially-consistent stores need a trailing fence; before that, a relaxed + # atomic is bracketed by (sequentially-consistent) fences + ordered = metal >= v"4.1" + args(order, flags, volatile) = air >= v"2.9" ? + "i32 $order, i32 2, i32 $flags, i1 $volatile" : "i32 $order, i32 2, i1 $volatile" + # loads and read-modify-writes are always volatile (see `select_atomic!`) + op(order; volatile=false) = + ordered ? args(order, order == 0 ? 0 : 3, volatile) : args(0, 0, true) + rmw(order) = op(order; volatile=true) + fence = "call void @air.atomic.fence(i32 3, i32 5, i32 2)" + ir = lower_metal_atomics(kernel(""" + %a = atomicrmw add ptr addrspace(1) %p, i32 1 seq_cst, align 4 + %b = load atomic i32, ptr addrspace(1) %p acquire, align 4 + store atomic i32 %b, ptr addrspace(1) %p release, align 4 + %c = cmpxchg ptr addrspace(1) %p, i32 1, i32 2 monotonic acquire, align 4 + %d = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4 + store atomic i32 %d, ptr addrspace(1) %p seq_cst, align 4 + """); metal, air) + @test @filecheck begin + @check_label "define void @f" + @check cond=!ordered fence + @check_next cond=!ordered "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, $(rmw(5)))" + @check cond=ordered "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, $(rmw(5)))" + @check_next cond=!ordered fence + @check cond=!ordered "call i32 @air.atomic.global.load.i32(ptr addrspace(1) %p, $(rmw(2)))" + @check_next cond=ordered "call i32 @air.atomic.global.load.i32(ptr addrspace(1) %p, $(rmw(2)))" + @check_next cond=!ordered fence + @check_next cond=!ordered fence + @check_next "call void @air.atomic.global.store.i32(ptr addrspace(1) %p, i32 {{%.+}}, $(op(3)))" + @check "call i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1) %p, ptr {{%.+}}, i32 2, i32 0, $(ordered ? args(2, 3, false) : args(0, 0, true)))" + @check_next "icmp eq i32 {{%.+}}, 1" + @check cond=!ordered fence + @check "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, $(rmw(0)))" + @check_next cond=!ordered fence + @check_next "call void @air.atomic.global.store.i32(ptr addrspace(1) %p, i32 {{%.+}}, $(op(5)))" + @check_next fence + @check_not "call void @air.atomic.fence" + ir + end + @test !occursin(r"(atomicrmw|cmpxchg|load atomic|store atomic) |^\s*fence "m, ir) + end + + @testset "relaxed loads (Metal $metal, AIR $air)" for (metal, air) in targets + # relaxed loads of device memory at device scope are acquire loads, so that they + # eventually see other threadgroups' stores; others aren't affected (and relaxed loads + # of device memory are selected at device scope regardless, like MSL does) + ir = lower_metal_atomics(kernel(""" + %a = load atomic i32, ptr addrspace(1) %p monotonic, align 4 + %b = load atomic i32, ptr addrspace(1) %p syncscope("workgroup") monotonic, align 4 + %c = load atomic i32, ptr addrspace(3) %t monotonic, align 4 + """); metal, air) + ordered = metal >= v"4.1" + args(order, flags) = air >= v"2.9" ? + "i32 $order, i32 2, i32 $flags, i1 true" : "i32 $order, i32 2, i1 true" + relaxed(scope) = air >= v"2.9" ? "i32 0, i32 $scope, i32 0, i1 true" : "i32 0, i32 $scope, i1 true" + @test @filecheck begin + @check "call i32 @air.atomic.global.load.i32(ptr addrspace(1) %p, $(ordered ? args(2, 3) : args(0, 0)))" + @check_next cond=!ordered "call void @air.atomic.fence(i32 3, i32 5, i32 2)" + @check_next "call i32 @air.atomic.global.load.i32(ptr addrspace(1) %p, $(relaxed(2)))" + @check_next "call i32 @air.atomic.local.load.i32(ptr addrspace(3) %t, $(relaxed(1)))" + ir + end + end + + @testset "sequentially-consistent stores" begin + # from MSL 4.1, only sequentially-consistent stores get a (trailing) fence, with the + # store's scope: other orderings and operations don't need one + ir = lower_metal_atomics(kernel(""" + store atomic i32 1, ptr addrspace(1) %p seq_cst, align 4 + store atomic i32 1, ptr addrspace(3) %t syncscope("workgroup") seq_cst, align 4 + store atomic float 1.0, ptr addrspace(1) %p seq_cst, align 4 + store atomic i32 1, ptr addrspace(1) %p release, align 4 + %a = load atomic i32, ptr addrspace(1) %p seq_cst, align 4 + %b = atomicrmw xchg ptr addrspace(1) %p, i32 1 seq_cst, align 4 + %c = cmpxchg ptr addrspace(1) %p, i32 1, i32 2 seq_cst seq_cst, align 4 + """); metal=v"4.1", air=v"2.9") + @test @filecheck begin + @check "call void @air.atomic.global.store.i32(ptr addrspace(1) %p, i32 1, i32 5, i32 2, i32 3, i1 false)" + @check_next "call void @air.atomic.fence(i32 3, i32 5, i32 2)" + @check_next "call void @air.atomic.local.store.i32(ptr addrspace(3) %t, i32 1, i32 5, i32 1, i32 3, i1 false)" + @check_next "call void @air.atomic.fence(i32 3, i32 5, i32 1)" + @check_next "call void @air.atomic.global.store.i32(ptr addrspace(1) %p, i32 1065353216, i32 5, i32 2, i32 3, i1 false)" + @check_next "call void @air.atomic.fence(i32 3, i32 5, i32 2)" + @check_next "call void @air.atomic.global.store.i32(ptr addrspace(1) %p, i32 1, i32 3, i32 2, i32 3, i1 false)" + @check_not "call void @air.atomic.fence" + ir + end + end + + @testset "scopes" begin + ir = lower_metal_atomics(kernel(""" + %a = atomicrmw add ptr addrspace(1) %p, i32 1 syncscope("singlethread") monotonic, align 4 + %b = atomicrmw add ptr addrspace(1) %p, i32 1 syncscope("subgroup") monotonic, align 4 + %c = atomicrmw add ptr addrspace(1) %p, i32 1 syncscope("workgroup") monotonic, align 4 + %d = atomicrmw add ptr addrspace(1) %p, i32 1 syncscope("device") monotonic, align 4 + %e = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4 + %f = atomicrmw add ptr addrspace(3) %t, i32 1 monotonic, align 4 + %g = atomicrmw add ptr addrspace(3) %t, i32 1 syncscope("device") monotonic, align 4 + %h = atomicrmw add ptr addrspace(3) %t, i32 1 syncscope("subgroup") monotonic, align 4 + fence syncscope("workgroup") release + fence syncscope("subgroup") acquire + fence seq_cst + """); metal=v"4.1", air=v"2.9") + @test @filecheck begin + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 0, i32 0, i1 true)" + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 4, i32 0, i1 true)" + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 1, i32 0, i1 true)" + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 true)" + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 true)" + @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 1, i32 0, i1 true)" + @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 1, i32 0, i1 true)" + @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 4, i32 0, i1 true)" + @check "call void @air.atomic.fence(i32 3, i32 3, i32 1)" + @check "call void @air.atomic.fence(i32 3, i32 2, i32 4)" + @check "call void @air.atomic.fence(i32 3, i32 5, i32 2)" + ir + end + end + + @testset "operations" begin + ir = lower_metal_atomics(kernel(""" + %xchg = atomicrmw xchg ptr addrspace(1) %p, i32 1 monotonic, align 4 + %add = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4 + %sub = atomicrmw sub ptr addrspace(1) %p, i32 1 monotonic, align 4 + %and = atomicrmw and ptr addrspace(1) %p, i32 1 monotonic, align 4 + %or = atomicrmw or ptr addrspace(1) %p, i32 1 monotonic, align 4 + %xor = atomicrmw xor ptr addrspace(1) %p, i32 1 monotonic, align 4 + %max = atomicrmw max ptr addrspace(1) %p, i32 1 monotonic, align 4 + %min = atomicrmw min ptr addrspace(1) %p, i32 1 monotonic, align 4 + %umax = atomicrmw umax ptr addrspace(1) %p, i32 1 monotonic, align 4 + %umin = atomicrmw umin ptr addrspace(1) %p, i32 1 monotonic, align 4 + %fadd = atomicrmw fadd ptr addrspace(1) %p, float 1.0 monotonic, align 4 + %fsub = atomicrmw fsub ptr addrspace(3) %t, float 1.0 monotonic, align 4 + store atomic i32 1, ptr addrspace(1) %p monotonic, align 4 + store atomic volatile i32 1, ptr addrspace(1) %p monotonic, align 4 + """); metal=v"4.1", air=v"2.9") + # loads and read-modify-writes are always volatile, stores only when LLVM's are + trailer = "i32 0, i32 2, i32 0, i1 true)" + @test @filecheck begin + @check "call i32 @air.atomic.global.xchg.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.sub.s.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.and.s.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.or.s.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.xor.s.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.max.s.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.min.s.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.max.u.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call i32 @air.atomic.global.min.u.i32(ptr addrspace(1) %p, i32 1, $trailer" + @check "call float @air.atomic.global.add.f32(ptr addrspace(1) %p, float 1.000000e+00, $trailer" + @check "call float @air.atomic.local.sub.f32(ptr addrspace(3) %t, float 1.000000e+00, i32 0, i32 1, i32 0, i1 true)" + @check "call void @air.atomic.global.store.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 false)" + @check "call void @air.atomic.global.store.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 true)" + # declared like Apple does, with element types for the typed-pointer downgrader + @check "declare !arg_eltypes [[I32:![0-9]+]] i32 @air.atomic.global.xchg.i32(ptr addrspace(1), i32, i32, i32, i32, i1) [[ATTRS:#[0-9]+]]" + @check "declare !arg_eltypes [[F32:![0-9]+]] float @air.atomic.global.add.f32" + @check "attributes [[ATTRS]] = { mustprogress nounwind willreturn }" + @check "[[I32]] = !{i32 0, i32 0}" + @check "[[F32]] = !{i32 0, float 0.000000e+00}" + ir + end + end + + @testset "compare-exchange" begin + # AIR's compare-exchange takes the expected value by reference and returns the old + # value; like MSL, the success flag compares that with the expected value + ir = lower_metal_atomics(kernel(""" + %pair = cmpxchg ptr addrspace(1) %p, i32 1, i32 2 acq_rel acquire, align 4 + %weak = cmpxchg weak ptr addrspace(3) %t, i32 1, i32 2 monotonic monotonic, align 4 + %ok = extractvalue { i32, i1 } %pair, 1 + %okw = extractvalue { i32, i1 } %weak, 1 + %both = and i1 %ok, %okw + store i1 %both, ptr addrspace(1) %p + """); metal=v"4.1", air=v"2.9") + @test @filecheck begin + @check "store i32 1, ptr [[EXP:%.+]], align 4" + @check_next "[[OLD:%.+]] = call i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1) %p, ptr [[EXP]], i32 2, i32 4, i32 2, i32 2, i32 3, i1 false)" + @check_next "icmp eq i32 [[OLD]], 1" + @check "call i32 @air.atomic.local.cmpxchg.weak.i32(ptr addrspace(3) %t, ptr {{%.+}}, i32 2, i32 0, i32 0, i32 1, i32 0, i1 false)" + @check "declare !arg_eltypes [[MD:![0-9]+]] i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1), ptr, i32, i32, i32, i32, i32, i1)" + @check "[[MD]] = !{i32 0, i32 0, i32 1, i32 0}" + ir + end + end + + @testset "expansions (Metal $metal, AIR $air)" for (metal, air) in targets + # operations AIR lacks become compare-exchange loops, 8- and 16-bit ones masked + # operations on the containing 32-bit word + ir = lower_metal_atomics(kernel(""" + %nand = atomicrmw nand ptr addrspace(1) %p, i32 1 monotonic, align 4 + %fmax = atomicrmw fmax ptr addrspace(1) %p, float 1.0 monotonic, align 4 + %fload = load atomic float, ptr addrspace(1) %p monotonic, align 4 + store atomic float %fload, ptr addrspace(3) %t monotonic, align 4 + %fxchg = atomicrmw xchg ptr addrspace(1) %p, float 1.0 monotonic, align 4 + %tgfadd = atomicrmw fadd ptr addrspace(3) %t, float 1.0 monotonic, align 4 + """); metal, air) + @test @filecheck begin + @check_label "define void @f" + # nand: a loop around a compare-exchange of the computed value + @check "atomic.global.load.i32" + @check "atomicrmw.start" + @check "and i32" + @check_next "xor i32 {{.+}}, -1" + @check "atomic.global.cmpxchg.weak.i32" + # fmax + @check "atomicrmw.start" + @check "call float @llvm.maxnum.f32" + @check "atomic.global.cmpxchg.weak.i32" + # floating-point loads, stores and exchanges on integers + @check "atomic.global.load.i32" + @check_next "bitcast i32 {{.+}} to float" + @check "bitcast float {{.+}} to i32" + @check_next "atomic.local.store.i32" + @check "atomic.global.xchg.i32(ptr addrspace(1) %p, i32 1065353216" + # threadgroup fadd needs MSL 4.1 + @check cond=(metal >= v"4.1") "call float @air.atomic.local.add.f32" + @check cond=(metal < v"4.1") "atomic.local.cmpxchg.weak.i32" + @check_not "atomicrmw {{[a-z]+}} ptr" + ir + end + + ir = lower_metal_atomics(kernel(""" + %add8 = atomicrmw add ptr addrspace(1) %p, i8 1 monotonic, align 1 + %or8 = atomicrmw or ptr addrspace(1) %p, i8 1 monotonic, align 1 + %l8 = load atomic i8, ptr addrspace(3) %t monotonic, align 1 + store atomic i16 1, ptr addrspace(1) %p monotonic, align 2 + %c16 = cmpxchg ptr addrspace(1) %p, i16 1, i16 2 monotonic monotonic, align 2 + %h = atomicrmw fadd ptr addrspace(1) %p, half 1.0 monotonic, align 2 + store atomic i8 1, ptr addrspace(3) %t unordered, align 1 + """); metal, air) + @test @filecheck begin + @check_label "define void @f" + # add: masked compare-exchange loop on the containing word + @check "ptrtoint ptr addrspace(1) %p to i64" + @check "and i64 {{.+}}, 3" + @check "atomicrmw.start" + @check "atomic.global.cmpxchg.weak.i32" + # or: a word-sized or, with the value shifted into place + @check "[[V:%.+]] = shl i32 1, {{%.+}}" + @check_next "call i32 @air.atomic.global.or.s.i32(ptr addrspace(1) {{%.+}}, i32 [[V]]," + # load: a word-sized load, shifted + @check "call i32 @air.atomic.local.load.i32" + @check_next "lshr i32" + @check_next "trunc i32 {{.+}} to i8" + # store and compare-exchange + @check "atomic.global.cmpxchg.weak.i32" + @check "partword.cmpxchg.loop" + @check "atomic.global.cmpxchg.weak.i32" + # half-precision add + @check "fadd half" + @check "atomic.global.cmpxchg.weak.i32" + # an unordered store, which becomes a (monotonic) compare-exchange loop + @check "atomic.local.cmpxchg.weak.i32" + @check_not "atomicrmw {{[a-z]+}} ptr" + ir + end + end + + @testset "64-bit min/max (Metal $metal, AIR $air)" for (metal, air) in targets + ir = lower_metal_atomics(kernel(""" + %a = atomicrmw umax ptr addrspace(1) %p, i64 1 monotonic, align 8 + %b = atomicrmw umin ptr addrspace(1) %p, i64 1 release, align 8 + """); metal, air) + trailer(order, flags, volatile) = air >= v"2.9" ? + "i32 $order, i32 2, i32 $flags, i1 $volatile" : "i32 $order, i32 2, i1 $volatile" + ordered = metal >= v"4.1" + @test @filecheck begin + @check "call void @air.atomic.global.max.u.i64(ptr addrspace(1) %p, i64 1, $(trailer(0, 0, true)))" + @check cond=!ordered "call void @air.atomic.fence(i32 3, i32 5, i32 2)" + @check "call void @air.atomic.global.min.u.i64(ptr addrspace(1) %p, i64 1, $(ordered ? trailer(3, 3, true) : trailer(0, 0, true)))" + ir + end + end + + @testset "intrinsics emitted by front-ends (Metal $metal, AIR $air)" for (metal, air) in targets + # in the MSL 4.1 form, rewritten for the target (as Metal.jl's did in `finish_ir!`) + ir = lower_metal_atomics(""" + declare i32 @air.atomic.global.add.s.i32(ptr addrspace(1), i32, i32, i32, i32, i1) + declare i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1), ptr, i32, i32, i32, i32, i32, i1) + define void @f(ptr addrspace(1) %p, ptr %e) { + %a = call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 false) + %b = call i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1) %p, ptr %e, i32 1, i32 0, i32 0, i32 2, i32 0, i1 false) + %c = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4 + ret void + } + """; metal, air) + trailer = metal >= v"4.1" ? "i32 0, i32 2, i32 0, i1 false" : + air >= v"2.9" ? "i32 0, i32 2, i32 0, i1 true" : "i32 0, i32 2, i1 true" + # (the read-modify-write we select is always volatile) + selected = air >= v"2.9" ? "i32 0, i32 2, i32 0, i1 true" : "i32 0, i32 2, i1 true" + @test @filecheck begin + @check "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, $trailer)" + @check "call i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1) %p, ptr %e, i32 1, i32 0, $trailer)" + @check "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, $selected)" + ir + end + @test occursin(r"declare !arg_eltypes ![0-9]+ i32 @air.atomic.global.cmpxchg.weak.i32", ir) + end + + # orderings and memory flags on intrinsics need MSL 4.1 + @test_throws "requires MSL 4.1" lower_metal_atomics(""" + declare i32 @air.atomic.global.add.s.i32(ptr addrspace(1), i32, i32, i32, i32, i1) + define void @f(ptr addrspace(1) %p) { + %a = call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 5, i32 2, i32 1, i1 false) + ret void + } + """; metal=v"4.0", air=v"2.8") + + @testset "thread-private memory" begin + # atomics on the thread's stack (e.g. from `AllocOpt`) become plain accesses, whatever + # their ordering, operation or size (GPUCompiler.jl#934) + ir = lower_metal_atomics(kernel(""" + %slot = alloca i64, align 8 + %hi = getelementptr inbounds i32, ptr %slot, i64 1 + store i32 0, ptr %slot, align 4 + %pair = cmpxchg ptr %slot, i32 0, i32 1 seq_cst monotonic, align 4 + %fadd = atomicrmw fadd ptr %hi, float 1.0 seq_cst, align 4 + %sel = select i1 %c, ptr %slot, ptr %hi + %umax = atomicrmw umax ptr %sel, i32 2 acquire, align 4 + %wide = load atomic i64, ptr %slot seq_cst, align 8 + store atomic volatile i64 %wide, ptr %slot release, align 8 + %old = extractvalue { i32, i1 } %pair, 0 + store i32 %old, ptr addrspace(1) %p + """; args="ptr addrspace(1) %p, i1 %c"); metal=v"3.2", air=v"2.7") + @test @filecheck begin + @check_label "define void @f" + # compare-exchange + @check "[[OLD:%.+]] = load i32, ptr %slot, align 4" + @check_next "[[OK:%.+]] = icmp eq i32 [[OLD]], 0" + @check_next "[[NEW:%.+]] = select i1 [[OK]], i32 1, i32 [[OLD]]" + @check_next "store i32 [[NEW]], ptr %slot, align 4" + # read-modify-write, also through a select of stack pointers + @check "load float, ptr %hi, align 4" + @check_next "fadd float" + @check_next "store float" + @check "load i32, ptr %sel, align 4" + @check_next "icmp ugt i32" + @check_next "select i1" + @check_next "store i32" + # 64-bit loads and stores + @check "load i64, ptr %slot, align 8" + @check_next "store volatile i64 {{%.+}}, ptr %slot, align 8" + @check_not "air.atomic" + ir + end + + # also through phis of stack pointers + ir = lower_metal_atomics(kernel(""" + %a = alloca i32, align 4 + %b = alloca i32, align 4 + br i1 %c, label %l, label %r + l: + br label %m + r: + br label %m + m: + %phi = phi ptr [ %a, %l ], [ %b, %r ] + %x = atomicrmw add ptr %phi, i32 1 monotonic, align 4 + """; args="i1 %c"); metal=v"4.1", air=v"2.9") + @test occursin("load i32, ptr %phi", ir) + @test !occursin("air.atomic", ir) + end + + @testset "unsupported atomics" begin + for (body, reason) in ( + ("%a = atomicrmw add ptr %g, i32 1 monotonic, align 4", + "atomic operation in address space 0"), + ("%a = atomicrmw add ptr addrspace(2) %c, i32 1 monotonic, align 4", + "atomic operation in address space 2"), + ("%a = atomicrmw add ptr addrspace(1) %p, i64 1 monotonic, align 8", + "64-bit atomic operation"), + ("%a = atomicrmw umax ptr addrspace(1) %p, i64 1 monotonic, align 8\nstore i64 %a, ptr addrspace(1) %p", + "64-bit atomic operation"), + ("%a = load atomic i64, ptr addrspace(1) %p monotonic, align 8", + "64-bit atomic operation"), + ("%a = load atomic ptr, ptr addrspace(1) %p acquire, align 8", + "64-bit atomic operation"), + ("%a = atomicrmw fadd ptr addrspace(1) %p, double 1.0 monotonic, align 8", + "64-bit atomic operation"), + ("%a = atomicrmw add ptr addrspace(3) %t, i64 1 monotonic, align 8", + "64-bit atomic operation"), + ("%a = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 2", + "misaligned atomic operation"), + ("%a = atomicrmw add ptr addrspace(1) %p, i32 1 syncscope(\"agent\") monotonic, align 4", + "synchronization scope \"agent\""), + # a pointer that may not be on the stack + ("%s = alloca i32, align 4\n%m = select i1 %b, ptr %s, ptr %g\n%a = atomicrmw add ptr %m, i32 1 monotonic, align 4", + "atomic operation in address space 0"), + ) + errors = lower_metal_atomics(kernel(body; args="ptr addrspace(1) %p, ptr addrspace(3) %t, ptr %g, ptr addrspace(2) %c, i1 %b"); + metal=v"4.1", air=v"2.9") + @test occursin(reason, errors) + end + # without fences, ordered atomics need MSL 4.1; with them, MSL 3.2 + @test occursin("ordered atomic operation", + lower_metal_atomics(kernel("%a = atomicrmw add ptr addrspace(1) %p, i32 1 seq_cst, align 4"); + metal=v"3.1", air=v"2.6")) + end + + # end-to-end, from Julia code emitting LLVM atomics (like UnsafeAtomics does), which also + # exercises typed pointers on Julia versions that still use them + mod = @eval module $(gensym()) + using LLVM, LLVM.Interop + @generated function atomic_rmw(ptr::Core.LLVMPtr{T,A}, val::T, ::Val{op}) where {T,A,op} + @dispose ctx=Context() begin + T_val = convert(LLVMType, T) + T_ptr = convert(LLVMType, ptr) + f, _ = create_function(T_val, [T_ptr, T_val]) + @dispose builder=IRBuilder() begin + position!(builder, BasicBlock(f, "entry")) + typed_ptr = bitcast!(builder, parameters(f)[1], LLVM.PointerType(T_val, A)) + rv = atomic_rmw!(builder, op, typed_ptr, parameters(f)[2], + LLVM.API.LLVMAtomicOrderingSequentiallyConsistent, + SyncScope("device")) + ret!(builder, rv) + end + call_function(f, T, Tuple{Core.LLVMPtr{T,A},T}, :ptr, :val) + end + end + function kernel(p::Core.LLVMPtr{Int32,1}, b::Core.LLVMPtr{UInt8,1}, + f::Core.LLVMPtr{Float32,1}) + x = atomic_rmw(p, Int32(1), Val(LLVM.API.LLVMAtomicRMWBinOpAdd)) + y = atomic_rmw(b, UInt8(1), Val(LLVM.API.LLVMAtomicRMWBinOpAdd)) + z = atomic_rmw(f, 1f0, Val(LLVM.API.LLVMAtomicRMWBinOpFMax)) + unsafe_store!(p, x + y + reinterpret(Int32, z)) + return + end + end + @testset "end-to-end (Metal $metal, AIR $air)" for (metal, air) in targets + source = methodinstance(typeof(mod.kernel), + Tuple{Core.LLVMPtr{Int32,1},Core.LLVMPtr{UInt8,1}, + Core.LLVMPtr{Float32,1}}, + Base.get_world_counter()) + target = MetalCompilerTarget(; macos=v"27", metal, air) + job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true)) + llvm = sprint(io -> GPUCompiler.code_llvm(io, job; dump_module=true)) + @test occursin("atomicrmw add", llvm) + @test occursin("atomicrmw fmax", llvm) + asm = sprint(io -> GPUCompiler.code_native(io, job; dump_module=true)) + @test occursin("@air.atomic.global.add.s.i32", asm) + @test occursin("@air.atomic.global.cmpxchg.weak.i32", asm) + @test occursin("@air.atomic.fence", asm) == (metal < v"4.1") + @test !occursin(r"(atomicrmw|cmpxchg|load atomic|store atomic) ", asm) + end +end + +@testset "atomics on thread-private objects" begin + # GPUCompiler.jl#934: `AllocOpt` moves a non-escaping object with `@atomic` fields to the + # stack, keeping its compare-exchange loops, which Metal cannot express. From Julia 1.13, + # `@atomic` modifications are calls to `julia.atomicmodify`, which `AllocOpt` treats as + # an escape, so the object stays on the heap (the "thread-private memory" testset covers + # the lowering on every version). + @static if VERSION < v"1.13-" + mod = @eval module $(gensym()) + mutable struct Acc + @atomic n::Int32 + @atomic x::Float32 + end + function kernel(out::Core.LLVMPtr{Float32,1}, x::Float32) + acc = Acc(0, 0f0) + @atomic acc.n += Int32(1) + @atomic acc.x += x + unsafe_store!(out, (@atomic acc.n) + (@atomic acc.x)) + return + end + end + source = methodinstance(typeof(mod.kernel), Tuple{Core.LLVMPtr{Float32,1}, Float32}, + Base.get_world_counter()) + target = MetalCompilerTarget(; macos=v"15", metal=v"3.2", air=v"2.7") + job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true)) + @test @filecheck begin + @check "alloca" + @check "cmpxchg" + GPUCompiler.code_llvm(job; dump_module=true) + end + @test @filecheck begin + @check_not "cmpxchg" + @check_not "air.atomic" + GPUCompiler.code_native(job; dump_module=true) + end end end @testset "fence lowering" begin orders = [("acquire", 2), ("release", 3), ("acq_rel", 4), ("seq_cst", 5)] scopes = [("", 2), ("syncscope(\"singlethread\") ", 0), - ("syncscope(\"workgroup\") ", 2), + ("syncscope(\"workgroup\") ", 1), (raw"syncscope(\"fence acquire\22\5C\0A\") ", 2)] metadata = ("", ", !dbg !3, !annotation !4") fences_ir = join(["fence $scope$order$md" for (order, _) in orders @@ -1978,6 +2501,18 @@ end end end + # scopes without an MSL equivalent are rejected, not guessed + Context() do ctx + mod = parse(LLVM.Module, """ + define void @f() { + fence syncscope("block") seq_cst + ret void + } + """) + errors = GPUCompiler.validate_ir(fence_job(v"4.1"), mod) + @test only(errors)[1] == "fence with synchronization scope \"block\"" + end + # end-to-end: Julia's `atomic_fence` intrinsic must not reach the AIR as a bare `fence` # (Julia 1.14 added a syncscope argument to the intrinsic, JuliaLang/julia#60311) function kernel(p::Core.LLVMPtr{Int,1})