From 9ff929af446dc38cb0373efb66e1191567237996 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 14:25:38 +0200 Subject: [PATCH 1/7] Metal: map fence synchronization scopes like MSL does. `lower_fences!` used thread scope for `singlethread` fences and device scope for all others. Map the scopes the LLVM SPIR-V back-end uses to their MSL equivalents instead: `subgroup` to simdgroup, `workgroup` to threadgroup, and `device` and the system scope to device (Metal has no wider scope). Reject other scopes in `validate_ir`, as the NVPTX and AMDGPU back-ends do, rather than guessing what they mean. The scope and memory-order mappings are shared with the atomics lowering that follows, as are the target-independent helpers in `atomics.jl`. --- src/GPUCompiler.jl | 1 + src/atomics.jl | 13 ++++++++ src/metal.jl | 75 +++++++++++++++++++++++++++++++--------------- test/metal.jl | 14 ++++++++- 4 files changed, 78 insertions(+), 25 deletions(-) create mode 100644 src/atomics.jl diff --git a/src/GPUCompiler.jl b/src/GPUCompiler.jl index b74253ca9..5599e2d1a 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 000000000..7e90a43ad --- /dev/null +++ b/src/atomics.jl @@ -0,0 +1,13 @@ +# 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 + +# 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 diff --git a/src/metal.jl b/src/metal.jl index 3180b02fd..87cd64eff 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -465,6 +465,13 @@ function validate_ir(job::CompilerJob{MetalCompilerTarget}, mod::LLVM.Module) # Metal never supports 128-bit integers append!(errors, check_ir_values(mod, LLVM.IntType(128))) + # fences with a synchronization scope `lower_fences!` cannot map + for f in functions(mod), bb in blocks(f), inst in instructions(bb) + inst isa LLVM.FenceInst && metal_thread_scope(inst) === nothing || continue + push!(errors, ("fence with synchronization scope $(syncscope_name(inst))", + backtrace(inst), string(inst))) + end + errors end @@ -708,6 +715,41 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L return end +# 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 + # 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 +766,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 +791,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) diff --git a/test/metal.jl b/test/metal.jl index d5d5af901..707da4f89 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1929,7 +1929,7 @@ 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 +1978,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}) From 7371bc2439e4a84605f32ea918efc25a246a524c Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 16:30:47 +0200 Subject: [PATCH 2/7] Metal: lower LLVM atomics to AIR atomic intrinsics. Metal has no LLVM back-end, so LLVM's atomic instructions reached Apple's compiler as-is. It accepts the common cases, but it cannot expand what AIR lacks (8- and 16-bit operations, nand, floating-point min/max), ignores synchronization scopes, and only has ordered atomics from MSL 4.1. Do what `AtomicExpand` and instruction selection do for other targets, so that front-ends can emit plain LLVM atomics with an ordering and a scope: - `validate_ir` rejects the atomics that cannot be lowered: outside device and threadgroup memory, wider than 32 bits (except non-fetching 64-bit umin/umax on device memory), misaligned, with an unknown scope, or ordered below MSL 3.2; - below MSL 4.1, ordered operations are relaxed ones bracketed with fences (AtomicExpand's default leading/trailing scheme); - floating-point loads, stores and exchanges are cast to integers, 8- and 16-bit operations become masked operations on the containing word, and operations AIR lacks become compare-exchange loops; - the rest are selected to `air.atomic.*` calls, in the form of the target's AIR and MSL versions. `cmpxchg` becomes a weak compare-exchange with the success flag derived from the old value, as MSL does. One legality function, `metal_atomic_action`, decides both: like the rule tables of LLVM's legalizers, it returns how to lower an operation, or why it cannot be, which `validate_ir` reports. Front-ends that call `air.atomic.*` directly (e.g. to pass memory flags) do so in the MSL 4.1 form, which is now legalized for older targets here instead of in Metal.jl. The intrinsics are declared like Apple does (`mustprogress nounwind willreturn`), no longer `argmemonly` and `readonly`/`writeonly`, which would allow LLVM to move other memory accesses across ordered atomics. --- src/atomics.jl | 97 ++++++++ src/metal.jl | 617 +++++++++++++++++++++++++++++++++++++++++++++++-- test/metal.jl | 381 +++++++++++++++++++++++++++++- 3 files changed, 1070 insertions(+), 25 deletions(-) diff --git a/src/atomics.jl b/src/atomics.jl index 7e90a43ad..984f14c64 100644 --- a/src/atomics.jl +++ b/src/atomics.jl @@ -6,8 +6,105 @@ 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 87cd64eff..df493167e 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -465,11 +465,17 @@ function validate_ir(job::CompilerJob{MetalCompilerTarget}, mod::LLVM.Module) # Metal never supports 128-bit integers append!(errors, check_ir_values(mod, LLVM.IntType(128))) - # fences with a synchronization scope `lower_fences!` cannot map + # atomics that `lower_atomics!` cannot lower for f in functions(mod), bb in blocks(f), inst in instructions(bb) - inst isa LLVM.FenceInst && metal_thread_scope(inst) === nothing || continue - push!(errors, ("fence with synchronization scope $(syncscope_name(inst))", - backtrace(inst), string(inst))) + 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 @@ -670,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) @@ -689,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 @@ -715,6 +724,28 @@ 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: +# +# - 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 : @@ -750,6 +781,560 @@ function metal_thread_scope(inst::LLVM.Instruction, as::Union{Nothing,Int}=nothi 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 + +# 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): +# +# - `: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 + + 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 + +# 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 + +# 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)) + scope = ConstantInt(T_i32, metal_thread_scope(inst, as)) + flags = ConstantInt(T_i32, is_ordered(order) ? METAL_MEM_FLAGS : 0) + volatile = ConstantInt(T_i1, target.metal < v"4.1" || is_volatile(inst)) + 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 + job.config.target.metal < v"4.1" && insert_atomic_fences!(inst) + 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) @@ -2569,19 +3154,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 707da4f89..689dd35c1 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,8 +1925,375 @@ 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 false)" + @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); + # 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" + rmw(order) = ordered ? args(order, order == 0 ? 0 : 3, false) : args(0, 0, 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 "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 {{%.+}}, $(rmw(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 {{%.+}}, $(rmw(5)))" + @check_next cond=!ordered 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 "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 false)" + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 4, i32 0, i1 false)" + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 1, i32 0, i1 false)" + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 false)" + @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 false)" + @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 1, i32 0, i1 false)" + @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 1, i32 0, i1 false)" + @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 4, i32 0, i1 false)" + @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 + %vol = atomicrmw volatile add ptr addrspace(1) %p, i32 1 monotonic, align 4 + """); metal=v"4.1", air=v"2.9") + trailer = "i32 0, i32 2, i32 0, i1 false)" + @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 false)" + @check "call i32 @air.atomic.global.add.s.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, !ordered)))" + @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, false) : 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" + @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, $trailer)" + 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 "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\""), + ) + errors = lower_metal_atomics(kernel(body; args="ptr addrspace(1) %p, ptr addrspace(3) %t, ptr %g, ptr addrspace(2) %c"); + 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 From 0d349e7e9690ff630e7c882a839558310a3200c8 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 16:30:47 +0200 Subject: [PATCH 3/7] Metal: lower atomics on thread-private memory to plain accesses. When `AllocOpt` moves an object with `@atomic` fields to the stack, it keeps the compare-exchange loops on its fields (#934). MSL has no atomics on thread memory, so the lowering rejected those, and before it Apple's compiler handled them inconsistently: macOS 15's silently produced wrong results. They don't need to be atomic, though: no thread can access another's stack, so plain accesses behave the same, whatever the ordering or scope. Lower atomics on pointers that are derived from allocas only to plain loads and stores. --- src/metal.jl | 82 +++++++++++++++++++++++++++++++++++++++++++++ test/metal.jl | 92 ++++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 173 insertions(+), 1 deletion(-) diff --git a/src/metal.jl b/src/metal.jl index df493167e..017d0968e 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -732,6 +732,7 @@ end # 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 @@ -801,10 +802,41 @@ function atomic_bits(T::LLVMType) 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 @@ -820,6 +852,8 @@ function metal_atomic_action(@nospecialize(job::CompilerJob{MetalCompilerTarget} 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)" @@ -897,6 +931,50 @@ function set_atomic!(inst::LLVM.Instruction, order::AtomicOrdering, scope::SyncS 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`). @@ -1313,6 +1391,10 @@ function lower_atomics!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), 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 job.config.target.metal < v"4.1" && insert_atomic_fences!(inst) if action === :cast inst = cast_atomic_to_int!(job, inst) diff --git a/test/metal.jl b/test/metal.jl index 689dd35c1..da249f0a0 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -2217,6 +2217,61 @@ end } """; 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", @@ -2239,8 +2294,11 @@ end "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"); + 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 @@ -2297,6 +2355,38 @@ end 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 + 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 + @testset "fence lowering" begin orders = [("acquire", 2), ("release", 3), ("acq_rel", 4), ("seq_cst", 5)] scopes = [("", 2), ("syncscope(\"singlethread\") ", 0), From 370a9620ebf8e88e37e1958a1254cfb8373c3ded Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Thu, 24 Sep 2026 18:15:18 +0200 Subject: [PATCH 4/7] Only test thread-private atomics end-to-end before Julia 1.13. From Julia 1.13, `@atomic` modifications of fields are calls to `julia.atomicmodify`, which AllocOpt treats as an escape, so the object in the test stays on the heap and there are no atomics on the stack to lower. The IR-level test of the lowering still covers every version. --- test/metal.jl | 57 ++++++++++++++++++++++++++++----------------------- 1 file changed, 31 insertions(+), 26 deletions(-) diff --git a/test/metal.jl b/test/metal.jl index da249f0a0..c193e66ef 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -2357,33 +2357,38 @@ 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 - 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 + # 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 - 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 From 96f93cc550e1e85426cf299585c7ddc23ef06844 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 11:19:05 +0200 Subject: [PATCH 5/7] Metal: always set the volatile bit on atomic loads and read-modify-writes. MSL 4.1 clears the volatile bit of atomics on non-`volatile` objects, and we did the same for non-volatile LLVM atomics. Apple's compiler then treats the atomic like a plain access: a relaxed load of memory the kernel doesn't write is hoisted into the uniform preamble (out of a spin loop), and a read-modify-write that leaves memory unchanged (adding or or-ing 0, whatever the ordering) becomes a load. On the M1 (macOS 27), a spin-wait on such a read-modify-write never observes the store another threadgroup makes, while the volatile one does. LLVM atomics allow neither transformation, so set the bit on loads and read-modify-writes, as MSL does on every atomic before 4.1. Stores and compare-exchanges compile the same either way, as does any read-modify-write that does change memory, and the bit costs nothing measurable. --- src/metal.jl | 9 ++++++++- test/metal.jl | 46 +++++++++++++++++++++++++++------------------- 2 files changed, 35 insertions(+), 20 deletions(-) diff --git a/src/metal.jl b/src/metal.jl index 017d0968e..9dc2653e8 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -1244,7 +1244,14 @@ function select_atomic!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), order = lowered_ordering(job, atomic_ordering(inst)) scope = ConstantInt(T_i32, metal_thread_scope(inst, as)) flags = ConstantInt(T_i32, is_ordered(order) ? METAL_MEM_FLAGS : 0) - volatile = ConstantInt(T_i1, target.metal < v"4.1" || is_volatile(inst)) + # 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))) diff --git a/test/metal.jl b/test/metal.jl index c193e66ef..9ec4fdda7 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1928,7 +1928,7 @@ end 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 false)" + @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 @@ -1970,7 +1970,10 @@ end 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" - rmw(order) = ordered ? args(order, order == 0 ? 0 : 3, false) : args(0, 0, true) + # 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 @@ -1989,13 +1992,13 @@ end @check "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 {{%.+}}, $(rmw(3)))" + @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 {{%.+}}, $(rmw(5)))" + @check_next "call void @air.atomic.global.store.i32(ptr addrspace(1) %p, i32 {{%.+}}, $(op(5)))" @check_next cond=!ordered fence @check_not "call void @air.atomic.fence" ir @@ -2018,14 +2021,14 @@ end 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 false)" - @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 4, i32 0, i1 false)" - @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 1, i32 0, i1 false)" - @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 false)" - @check "@air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 false)" - @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 1, i32 0, i1 false)" - @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 1, i32 0, i1 false)" - @check "@air.atomic.local.add.s.i32(ptr addrspace(3) %t, i32 1, i32 0, i32 4, i32 0, i1 false)" + @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)" @@ -2047,9 +2050,11 @@ end %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 - %vol = atomicrmw volatile add ptr addrspace(1) %p, i32 1 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") - trailer = "i32 0, i32 2, i32 0, i1 false)" + # 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" @@ -2062,8 +2067,9 @@ end @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 false)" - @check "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 true)" + @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" @@ -2178,9 +2184,9 @@ end "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, !ordered)))" + @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, false) : trailer(0, 0, true)))" + @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 @@ -2199,10 +2205,12 @@ end """; 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, $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) From 10a594405084079f38bd84a2e8a485c4a0d0ab40 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 11:20:38 +0200 Subject: [PATCH 6/7] Metal: make sequentially-consistent stores sequentially consistent. From MSL 4.1, we select an LLVM seq_cst atomic to the MSL 4.1 seq_cst intrinsic, but Apple's compiler emits a seq_cst store as a release store that doesn't wait for its write. A later seq_cst load can then be performed before it: in the store-buffering litmus test (`store x; load y` against `store y; load x`, across threadgroups), both loads return the old values about 0.5% of the time on an M1, for LLVM atomics and MSL's own seq_cst atomics alike. Follow seq_cst stores by a seq_cst fence, like AtomicExpand does for targets whose `shouldInsertTrailingSeqCstFenceForAtomicStore` returns true. Read-modify-writes and compare-exchanges wait for their result and already behave; below MSL 4.1, the fences that bracket every ordered operation already did. On the M1, this removes all store-buffering violations (none in 6.5M runs, also for the RMW and compare-exchange variants), and costs 4-11% on a seq_cst-store-bound kernel. Fences before seq_cst loads work too, but cost 50% on a load-bound one. --- src/metal.jl | 26 +++++++++++++++++++++++++- test/metal.jl | 35 +++++++++++++++++++++++++++++++---- 2 files changed, 56 insertions(+), 5 deletions(-) diff --git a/src/metal.jl b/src/metal.jl index 9dc2653e8..b6338b211 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -1207,6 +1207,26 @@ function insert_atomic_fences!(inst::LLVM.Instruction) 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) = @@ -1402,7 +1422,11 @@ function lower_atomics!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), demote_private_atomic!(inst) continue end - job.config.target.metal < v"4.1" && insert_atomic_fences!(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 diff --git a/test/metal.jl b/test/metal.jl index 9ec4fdda7..6ea9b57ab 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -1965,8 +1965,9 @@ end 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); - # before that, a relaxed atomic is bracketed by (sequentially-consistent) fences + # 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" @@ -1989,7 +1990,8 @@ end @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 "call i32 @air.atomic.global.load.i32(ptr addrspace(1) %p, $(rmw(2)))" + @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)))" @@ -1999,13 +2001,38 @@ end @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 cond=!ordered fence + @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 "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 From faea1162532d145d2ad7c9b5676a59a7528f7bcc Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 25 Sep 2026 11:22:20 +0200 Subject: [PATCH 7/7] Metal: make relaxed loads of device memory eventually see other stores. MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit LLVM requires that a thread repeatedly loading an address monotonically eventually sees the stores of other threads. MSL doesn't: device memory is only coherent within a threadgroup by default, and relaxed atomics only guarantee atomicity and a consistent modification order (MSL 4.1 §4.8, §6.16.1.1). Apple's compiler emits relaxed atomic loads of device memory as cached loads, so on an M1 a spin loop that relaxed-loads a flag another threadgroup sets never sees the store: not in 20M iterations, on every MSL version, with MSL's own relaxed atomics too. (Within a threadgroup it does, as the threadgroup shares the cache.) Acquire loads invalidate the cache after loading, and the same spin loop exits after a few hundred iterations. Lower device-scope monotonic loads of device memory as acquire loads, which below MSL 4.1 means a trailing fence. Loads at a narrower scope only need to see stores from their threadgroup, so they are unaffected; front-ends that want MSL's semantics (and cost) can use the workgroup scope. Relaxed loads of device memory are selected at device scope regardless, like MSL does: the scope doesn't change the code, but it is recorded. This costs nothing when the loads miss the cache, but a kernel of cache-resident device-scope relaxed loads runs 6x (MSL 4.1) to 7x (older) slower. Only strengthening loads that may repeat on the same address avoids part of that, but needs cycle analysis that isn't worth its code. --- src/metal.jl | 24 +++++++++++++++++++++++- test/metal.jl | 22 ++++++++++++++++++++++ 2 files changed, 45 insertions(+), 1 deletion(-) diff --git a/src/metal.jl b/src/metal.jl index b6338b211..5c9515c6d 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -1207,6 +1207,23 @@ function insert_atomic_fences!(inst::LLVM.Instruction) 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 @@ -1262,7 +1279,11 @@ function select_atomic!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), # the operands after the memory order(s): scope, flags (from AIR 2.9), volatile order = lowered_ordering(job, atomic_ordering(inst)) - scope = ConstantInt(T_i32, metal_thread_scope(inst, as)) + # (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., @@ -1422,6 +1443,7 @@ function lower_atomics!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), demote_private_atomic!(inst) continue end + strengthen_relaxed_load!(job, inst) if job.config.target.metal < v"4.1" insert_atomic_fences!(inst) else diff --git a/test/metal.jl b/test/metal.jl index 6ea9b57ab..baf78780a 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -2008,6 +2008,28 @@ 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