Skip to content

Metal: stop legalizing intrinsics that front-ends emit #1047

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
maleadt merged 1 commit into from
Oct 21, 2026
Merged

Metal: stop legalizing intrinsics that front-ends emit #1047

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
136 changes: 9 additions & 127 deletions src/metal.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -729,8 +729,6 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
# the macOS 27 back-end rejects bare LLVM fences (Metal.jl#968)
lower_fences!(job, mod)

legalize_barrier_abi!(job, mod)

# strip device-side `trap`s and rewrite `unreachable` into clean returns (#433, #370). this
# runs post-`optimize!`, after the trap has finished serving as the optimizer guard; the pass
# force-inlines throwing functions into the kernel first so the rewrite is sound, then scrubs
Expand Down Expand Up @@ -828,10 +826,8 @@ end
# form the target's AIR and MSL versions use.
#
# Front-ends pass MSL's memory flags, which LLVM cannot express, in the synchronization scope
# (see `metal_syncscope`). They can also call `air.atomic.*` intrinsics directly, 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.
# (see `metal_syncscope`). `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::LLVM.AtomicOrdering.T) =
Expand Down Expand Up @@ -1169,27 +1165,12 @@ function select_atomic!(@nospecialize(job::CompilerJob{MetalCompilerTarget}),
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(mod.functions)
fn = f.name
m = match(r"^air\.atomic\.(global|local)\..+\.(i32|f32|i64)$", fn)
# Tell the downgrader the element types of the pointer operands of the `air.atomic.*`
# intrinsics, which LLVM cannot infer.
function annotate_atomic_eltypes!(mod::LLVM.Module)
for f in mod.functions
m = match(r"^air\.atomic\.(global|local)\..+\.(i32|f32|i64)$", f.name)
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(f.function_type.parameters)
Expand All @@ -1198,66 +1179,14 @@ function legalize_atomic_abi!(@nospecialize(job::CompilerJob{MetalCompilerTarget
push!(mds, null(T))
end
f.metadata["arg_eltypes"] = MDNode(mds)

target.metal >= v"4.1" && continue
nparams = length(f.function_type.parameters)
# older front-ends rewrote their calls for the target themselves
nparams == air_atomic_arity(fn) || continue

is_cmpxchg = occursin(".cmpxchg.", fn)
calls = [u::LLVM.CallInst for u in f.users]
for call in calls
args = collect(call.arguments)
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
call.operands[end-1] = ConstantInt(true)
end
continue
end

# before AIR 2.9, there is no flags operand
ft = f.function_type
params = ft.parameters
legacy_ft = LLVM.FunctionType(ft.return_type, [params[1:end-2]..., params[end]])
f.name = fn * ".msl41"
legacy_f = LLVM.Function(mod, fn, legacy_ft)
append!(legacy_f.function_attributes, f.function_attributes)
legacy_f.metadata["arg_eltypes"] = f.metadata["arg_eltypes"]
for call in calls
args = collect(call.arguments)
@dispose builder=IRBuilder() begin
position!(builder, LLVM.before(call))
builder.debug_location = call.debug_location
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 mod.functions for bb in f.blocks for inst in bb.instructions
if is_atomic_memop(inst)]
isempty(atomics) && return changed
isempty(atomics) && return false

for inst in atomics
action = metal_atomic_action(job, inst)
Expand Down Expand Up @@ -1302,7 +1231,7 @@ function lower_atomics!(@nospecialize(job::CompilerJob{MetalCompilerTarget}),
isempty(unsupported) || throw(InvalidIRError(job, unsupported))

# attach element type metadata to the declarations we introduced
legalize_atomic_abi!(job, mod)
annotate_atomic_eltypes!(mod)
return true
end

Expand Down Expand Up @@ -1347,53 +1276,6 @@ function lower_fences!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod
return true
end

# Rewrite the `air.wg.barrier` and `air.simdgroup.barrier` calls front-ends emit, in either
# form, into the form of the target: AIR 2.9 added a memory order operand between the memory
# flags and the scope, and Apple's tools reject the form of the other AIR versions. Barriers
# are sequentially consistent before MSL 4.1, so other orders are strengthened then.
function legalize_barrier_abi!(@nospecialize(job::CompilerJob{MetalCompilerTarget}),
mod::LLVM.Module)
target = job.config.target
T_int32 = LLVM.Int32Type()
seq_cst = ConstantInt(T_int32, 5)
changed = false
for fn in ("air.wg.barrier", "air.simdgroup.barrier")
haskey(mod.functions, fn) || continue
f = mod.functions[fn]
calls = [u::LLVM.CallInst for u in f.users]
nparams = length(f.function_type.parameters)
if target.air >= v"2.9" && nparams == 3
target.metal >= v"4.1" && continue
for call in calls
call.operands[2] = seq_cst
end
changed |= !isempty(calls)
continue
elseif target.air < v"2.9" && nparams == 2
continue
end

new_ft = LLVM.FunctionType(LLVM.VoidType(),
fill(T_int32, target.air >= v"2.9" ? 3 : 2))
f.name = ""
new_f = LLVM.Function(mod, fn, new_ft)
append!(new_f.function_attributes, f.function_attributes)
for call in calls
flags, scope = call.arguments[1], call.arguments[end]
args = target.air >= v"2.9" ? [flags, seq_cst, scope] : [flags, scope]
@dispose builder=IRBuilder() begin
position!(builder, LLVM.before(call))
builder.debug_location = call.debug_location
call!(builder, new_ft, new_f, args)
end
erase!(call)
end
erase!(f)
changed = true
end
return changed
end

# Julia names each generated LLVM function `julia_<name>_<counter>`, where the counter is
# drawn from a process-global codegen sequence and so differs from one session to the next.
# Left anywhere in the emitted bitcode it makes the AIR (and the metallib wrapping it)
Expand Down
128 changes: 0 additions & 128 deletions test/metal.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2947,40 +2947,6 @@ end
end
end

@testset "intrinsics emitted by front-ends (Metal $metal, AIR $air)" for (metal, air) in targets
# in the MSL 4.1 form, rewritten for the target (as Metal.jl's did in `finish_ir!`)
ir = lower_metal_atomics("""
declare i32 @air.atomic.global.add.s.i32(ptr addrspace(1), i32, i32, i32, i32, i1)
declare i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1), ptr, i32, i32, i32, i32, i32, i1)
define void @f(ptr addrspace(1) %p, ptr %e) {
%a = call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 0, i32 2, i32 0, i1 false)
%b = call i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1) %p, ptr %e, i32 1, i32 0, i32 0, i32 2, i32 0, i1 false)
%c = atomicrmw add ptr addrspace(1) %p, i32 1 monotonic, align 4
ret void
}
"""; metal, air)
trailer = metal >= v"4.1" ? "i32 0, i32 2, i32 0, i1 false" :
air >= v"2.9" ? "i32 0, i32 2, i32 0, i1 true" : "i32 0, i32 2, i1 true"
# (the read-modify-write we select is always volatile)
selected = air >= v"2.9" ? "i32 0, i32 2, i32 0, i1 true" : "i32 0, i32 2, i1 true"
@test @filecheck begin
@check "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, $trailer)"
@check "call i32 @air.atomic.global.cmpxchg.weak.i32(ptr addrspace(1) %p, ptr %e, i32 1, i32 0, $trailer)"
@check "call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, $selected)"
ir
end
@test occursin(r"declare !arg_eltypes ![0-9]+ i32 @air.atomic.global.cmpxchg.weak.i32", ir)
end

# orderings and memory flags on intrinsics need MSL 4.1
@test_throws "requires MSL 4.1" lower_metal_atomics("""
declare i32 @air.atomic.global.add.s.i32(ptr addrspace(1), i32, i32, i32, i32, i1)
define void @f(ptr addrspace(1) %p) {
%a = call i32 @air.atomic.global.add.s.i32(ptr addrspace(1) %p, i32 1, i32 5, i32 2, i32 1, i1 false)
ret void
}
"""; metal=v"4.0", air=v"2.8")

@testset "thread-private memory" begin
# atomics on the thread's stack (e.g. from `AllocOpt`) become plain accesses, whatever
# their ordering, operation or size (GPUCompiler.jl#934)
Expand Down Expand Up @@ -3286,100 +3252,6 @@ end
@test occursin("call void @air.atomic.fence(i32 3, i32 3, i32 2)", air)
end

@testset "barrier lowering" begin
# MSL's barriers as Apple's front-end emits them, before and after AIR 2.9
ir(arity) = """
declare void @air.wg.barrier($(join(fill("i32", arity), ", "))) #0
declare void @air.simdgroup.barrier($(join(fill("i32", arity), ", "))) #0
define void @f() {
call void @air.wg.barrier(i32 2, $(arity == 3 ? "i32 2, " : "")i32 1)
call void @air.simdgroup.barrier(i32 1, $(arity == 3 ? "i32 4, " : "")i32 4)
ret void
}
attributes #0 = { convergent nounwind }
"""
barrier_job(metal, air) = let
source = methodinstance(typeof(identity), Tuple{Int}, Base.get_world_counter())
target = MetalCompilerTarget(; macos=v"27", metal, air)
CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true))
end

@testset "Metal $metal, AIR $air, $arity operands" for
(metal, air) in ((v"3.2", v"2.7"), (v"4.0", v"2.8"), (v"4.0", v"2.9"), (v"4.1", v"2.9")),
arity in (2, 3)
Context() do ctx
mod = parse(LLVM.Module, ir(arity))
job = barrier_job(metal, air)
GPUCompiler.legalize_barrier_abi!(job, mod)
@test (verify(mod); true)
str = string(mod)
calls = [m.match for m in eachmatch(r"call void @air\.[a-z.]+\([^\n]+\)", str)]
if air >= v"2.9"
# an order the front-end chose is kept from MSL 4.1 on
wg, simd = metal >= v"4.1" && arity == 3 ? (2, 4) : (5, 5)
@test calls == ["call void @air.wg.barrier(i32 2, i32 $wg, i32 1)",
"call void @air.simdgroup.barrier(i32 1, i32 $simd, i32 4)"]
@test occursin(r"declare void @air.wg.barrier\(i32, i32, i32\) #0", str)
else
@test calls == ["call void @air.wg.barrier(i32 2, i32 1)",
"call void @air.simdgroup.barrier(i32 1, i32 4)"]
@test occursin(r"declare void @air.wg.barrier\(i32, i32\) #0", str)
end
@test occursin(r"attributes #0 = \{ convergent nounwind \}", str)
@test count("declare void @air.", str) == 2
end
end

# end-to-end, with the form Metal.jl emits, and checked by Apple's tools when available
function kernel(p::Core.LLVMPtr{Float32,1})
unsafe_store!(p, 1f0)
ccall("extern air.wg.barrier", llvmcall, Cvoid, (UInt32, UInt32), 2, 1)
unsafe_store!(p, 2f0)
return
end
function plain_kernel(p::Core.LLVMPtr{Float32,1})
unsafe_store!(p, 1f0)
return
end
metal_opt = something(Sys.which("metal-opt"),
Some(try readchomp(pipeline(`xcrun --find metal-opt`; stderr=devnull))
catch; nothing end))
function metal_opt_accepts(job; verbose=true)
obj = JuliaContext() do _
first(GPUCompiler.compile(:obj, job))
end
mktemp() do path, io
write(io, obj)
close(io)
log = IOBuffer()
ok = success(pipeline(`$metal_opt $path -o /dev/null`; stdout=log, stderr=log))
ok || !verbose || print(String(take!(log)))
ok
end
end
@testset "AIR $air" for (metal, air, macos) in ((v"4.0", v"2.8", v"26"),
(v"4.1", v"2.9", v"27"))
kernel_job(f) = let
source = methodinstance(typeof(f), Tuple{Core.LLVMPtr{Float32,1}},
Base.get_world_counter())
target = MetalCompilerTarget(; macos, metal, air)
CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true))
end
job = kernel_job(kernel)

asm = sprint(io->GPUCompiler.code_native(io, job; dump_module=true))
@test occursin(air >= v"2.9" ? "call void @air.wg.barrier(i32 2, i32 5, i32 1)" :
"call void @air.wg.barrier(i32 2, i32 1)", asm)

# (an older toolchain may not know the target's AIR version at all)
if metal_opt === nothing || !metal_opt_accepts(kernel_job(plain_kernel); verbose=false)
@test_skip false
else
@test metal_opt_accepts(job)
end
end
end

@testset "floating-point atomics on Apple10" begin
# the macOS 26 compiler for Apple10 GPUs crashes on a floating-point add or subtract on
# device memory that a release fence precedes, like the fences that implement ordered
Expand Down
Loading