Skip to content

Metal: legalize the simdgroup matrix and tensor intrinsics #1046

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

Closed
maleadt wants to merge 1 commit into from
Closed

Metal: legalize the simdgroup matrix and tensor intrinsics #1046

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: 136 additions & 0 deletions src/metal.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -730,6 +730,8 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
lower_fences!(job, mod)

legalize_barrier_abi!(job, mod)
legalize_simdgroup_matrix_abi!(job, mod)
legalize_tensor_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
Expand Down Expand Up @@ -1394,6 +1396,140 @@ function legalize_barrier_abi!(@nospecialize(job::CompilerJob{MetalCompilerTarge
return changed
end

# The statically-known integer in lane `i` of a vector value, or `nothing` if that lane isn't
# a compile-time constant. Looks through the `insertelement` chains and constant vectors that
# front-ends build vector operands from.
function static_vector_lane(v::LLVM.Value, i::Integer)
if v isa LLVM.InsertElementInst
base, elt, idx = v.operands
idx isa LLVM.ConstantInt || return nothing # unknown insert position
if convert(Int, idx) == i
return elt isa LLVM.ConstantInt ? convert(Int, elt) : nothing
end
return static_vector_lane(base, i) # this lane is untouched by the insert
end
# (lanes of other values, including undef and poison vectors, aren't known integers)
v isa Union{LLVM.ConstantVector, LLVM.ConstantDataVector,
LLVM.ConstantAggregateZero} || return nothing
el = get(v.elements, i+1, nothing)
return el isa LLVM.ConstantInt ? convert(Int, el) : nothing
end

# Redeclare `f` with type `ft`, under the same name, and rewrite its calls with `rewrite`,
# which is given the builder (positioned before the call) and the call's arguments, and
# returns the arguments of the new call.
function redeclare_intrinsic!(rewrite, mod::LLVM.Module, f::LLVM.Function,
ft::LLVM.FunctionType)
fn = f.name
calls = [u::LLVM.CallInst for u in f.users]
f.name = ""
new_f = LLVM.Function(mod, fn, ft)
append!(new_f.function_attributes, f.function_attributes)
for call in calls
@dispose builder=IRBuilder() begin
position!(builder, LLVM.before(call))
builder.debug_location = call.debug_location
new = call!(builder, ft, new_f, rewrite(builder, collect(call.arguments)))
replace_uses!(call, new)
end
erase!(call)
end
erase!(f)
return new_f
end

# Rewrite the calls to `air.simdgroup_matrix_8x8_{load,store}` that front-ends emit, in the
# form of AIR 2.8, into the form of the target. AIR 2.8 replaced the elements-per-row scalar,
# matrix origin and transposition flag with vectors describing the dimensions, strides and
# origin of the memory operand, with transposition expressed by swapping their elements.
# Only the legacy form is derived; calls already in it, as older front-ends emitted them, are
# left alone. Also tell the downgrader the element type of the pointer operand.
function legalize_simdgroup_matrix_abi!(@nospecialize(job::CompilerJob{MetalCompilerTarget}),
mod::LLVM.Module)
T_i64 = LLVM.Int64Type()
T_vec2 = LLVM.VectorType(T_i64, 2)
T_bool = LLVM.Int1Type()
changed = false
for f in collect(mod.functions)
fn = f.name
m = match(r"^air\.simdgroup_matrix_8x8_(load|store)\.v64(\w+)\.p\d+\w+$", fn)
m === nothing && continue
is_load = m.captures[1] == "load"
ptr_idx = is_load ? 1 : 2

# pointer element type for the typed-pointer downgrader
suffix = m.captures[2]
T = suffix == "f16" ? LLVM.HalfType() :
suffix == "bf16" ? LLVM.BFloatType() :
suffix == "f32" ? LLVM.FloatType() :
LLVM.IntType(parse(Int, suffix[2:end]))
eltypes = MDNode([ConstantInt(Int32(ptr_idx - 1)), null(T)])
f.metadata["arg_eltypes"] = eltypes

job.config.target.air >= v"2.8" && continue
params = f.function_type.parameters
# older front-ends emitted the legacy form themselves
last(params) == T_bool && continue

# elements per row; origin; transpose
legacy_ft = LLVM.FunctionType(f.function_type.return_type,
[params[1:ptr_idx]..., T_i64, T_vec2, T_bool])
changed |= !isempty(f.uses)
legacy_f = redeclare_intrinsic!(mod, f, legacy_ft) do builder, args
_, strides, origin = args[end-2:end]

# the transposed layout has a unit stride in the second dimension, the other one
# in the first. the legacy intrinsic needs the flag as an immediate, so one of the
# two has to be the literal 1.
transposed = static_vector_lane(strides, 1) == 1
transposed || static_vector_lane(strides, 0) == 1 ||
error("$fn with strides that are not statically unit in either " *
"dimension requires AIR 2.8; the target is AIR $(job.config.target.air)")

# elements per row is the non-unit (leading-dimension) stride
epr = extract_element!(builder, strides, ConstantInt(Int32(transposed ? 0 : 1)))

# the transposed layout swaps the origin's elements; the legacy form doesn't
if transposed
o1 = extract_element!(builder, origin, ConstantInt(Int32(0)))
o2 = extract_element!(builder, origin, ConstantInt(Int32(1)))
origin = insert_element!(builder, UndefValue(T_vec2), o2, ConstantInt(Int32(0)))
origin = insert_element!(builder, origin, o1, ConstantInt(Int32(1)))
end

[args[1:ptr_idx]..., epr, origin, ConstantInt(T_bool, transposed)]
end
legacy_f.metadata["arg_eltypes"] = eltypes
end
return changed
end

# Rewrite the calls to the tensor intrinsics that front-ends emit, in the form of AIR 2.9,
# into the form of the target. AIR 2.9 appended operands, which Apple's front-end always
# passes as constants (`null, 0, -1` to `air.init_strided_private_tensor`, presumably
# describing blockwise-quantization scales, and `0` to `air.get_extent_private_tensor`), so
# they are dropped before AIR 2.9.
function legalize_tensor_abi!(@nospecialize(job::CompilerJob{MetalCompilerTarget}),
mod::LLVM.Module)
job.config.target.air >= v"2.9" && return false
changed = false
for (fn, nparams) in ("air.init_strided_private_tensor.i32.global" => 6,
"air.init_strided_private_tensor.i32.local" => 6,
"air.get_extent_private_tensor.i32" => 3)
f = get(mod.functions, fn, nothing)
f === nothing && continue
ft = f.function_type
params = ft.parameters
# older front-ends emitted the AIR 2.8 form
length(params) == nparams && continue

changed |= !isempty(f.uses)
legacy_ft = LLVM.FunctionType(ft.return_type, params[1:nparams])
redeclare_intrinsic!((_, args) -> args[1:nparams], mod, f, legacy_ft)
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
126 changes: 126 additions & 0 deletions test/metal.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3380,6 +3380,132 @@ end
end
end

@testset "simdgroup matrix lowering" begin
# the AIR 2.8 form Metal.jl emits: a transposed load, and a store that isn't transposed
v2(a, b) = (VecElement(Int64(a)), VecElement(Int64(b)))
function kernel(p::Core.LLVMPtr{Float32,1})
m = ccall("extern air.simdgroup_matrix_8x8_load.v64f32.p1f32", llvmcall,
NTuple{64, VecElement{Float32}}, (Core.LLVMPtr{Float32,1}, NTuple{2, VecElement{Int64}}, NTuple{2, VecElement{Int64}},
NTuple{2, VecElement{Int64}}),
p, v2(8, 16), v2(16, 1), v2(3, 2))
ccall("extern air.simdgroup_matrix_8x8_store.v64f32.p1f32", llvmcall,
Cvoid, (NTuple{64, VecElement{Float32}}, Core.LLVMPtr{Float32,1}, NTuple{2, VecElement{Int64}}, NTuple{2, VecElement{Int64}},
NTuple{2, VecElement{Int64}}),
m, p, v2(16, 8), v2(1, 16), v2(2, 3))
return
end
# the legacy form, as older versions of Metal.jl emitted it for older targets
function legacy_kernel(p::Core.LLVMPtr{Float32,1})
m = ccall("extern air.simdgroup_matrix_8x8_load.v64f32.p1f32", llvmcall,
NTuple{64, VecElement{Float32}}, (Core.LLVMPtr{Float32,1}, Int64, NTuple{2, VecElement{Int64}}, Bool),
p, 16, v2(2, 3), true)
ccall("extern air.simdgroup_matrix_8x8_store.v64f32.p1f32", llvmcall,
Cvoid, (NTuple{64, VecElement{Float32}}, Core.LLVMPtr{Float32,1}, Int64, NTuple{2, VecElement{Int64}}, Bool),
m, p, 16, v2(2, 3), false)
return
end
function air_module(f, air)
source = methodinstance(typeof(f), Tuple{Core.LLVMPtr{Float32,1}},
Base.get_world_counter())
target = MetalCompilerTarget(; macos=v"26", metal=v"3.2", air)
job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true))
sprint(io->GPUCompiler.code_native(io, job; dump_module=true))
end
ptr = raw"(?:(?:float|\{\}) addrspace\(1\)\*|ptr addrspace\(1\))"
v64 = raw"<64 x float>"

@testset "AIR $air" for air in (v"2.7", v"2.8")
asm = air_module(kernel, air)
if air >= v"2.8"
@test occursin(Regex("call $v64 @air\\.simdgroup_matrix_8x8_load\\.v64f32\\.p1f32\\($ptr [^,]+, <2 x i64> <i64 8, i64 16>, <2 x i64> <i64 16, i64 1>, <2 x i64> <i64 3, i64 2>\\)"), asm)
@test occursin(Regex("call void @air\\.simdgroup_matrix_8x8_store\\.v64f32\\.p1f32\\($v64 [^,]+, $ptr [^,]+, <2 x i64> <i64 16, i64 8>, <2 x i64> <i64 1, i64 16>, <2 x i64> <i64 2, i64 3>\\)"), asm)
else
# the transposed origin is swapped back
@test occursin(Regex("call $v64 @air\\.simdgroup_matrix_8x8_load\\.v64f32\\.p1f32\\($ptr [^,]+, i64 16, <2 x i64> <i64 2, i64 3>, i1 true\\)"), asm)
@test occursin(Regex("call void @air\\.simdgroup_matrix_8x8_store\\.v64f32\\.p1f32\\($v64 [^,]+, $ptr [^,]+, i64 16, <2 x i64> <i64 2, i64 3>, i1 false\\)"), asm)
end
end

# with a leading dimension and origin only known at run time
function dynamic_kernel(p::Core.LLVMPtr{Float32,1}, n::Int64)
m = ccall("extern air.simdgroup_matrix_8x8_load.v64f32.p1f32", llvmcall,
NTuple{64, VecElement{Float32}},
(Core.LLVMPtr{Float32,1}, NTuple{2, VecElement{Int64}},
NTuple{2, VecElement{Int64}}, NTuple{2, VecElement{Int64}}),
p, v2(8, n), v2(n, 1), v2(n, 2))
ccall("extern air.simdgroup_matrix_8x8_store.v64f32.p1f32", llvmcall,
Cvoid, (NTuple{64, VecElement{Float32}}, Core.LLVMPtr{Float32,1},
NTuple{2, VecElement{Int64}}, NTuple{2, VecElement{Int64}},
NTuple{2, VecElement{Int64}}),
m, p, v2(8, n), v2(n, 1), v2(n, 2))
return
end
let source = methodinstance(typeof(dynamic_kernel), Tuple{Core.LLVMPtr{Float32,1}, Int64},
Base.get_world_counter())
target = MetalCompilerTarget(; macos=v"26", metal=v"3.2", air=v"2.7")
job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true))
asm = sprint(io->GPUCompiler.code_native(io, job; dump_module=true))
@test occursin(r"@air\.simdgroup_matrix_8x8_load\.v64f32\.p1f32\([^,]+ [^,]+, i64 %[^,]+, <2 x i64> %[^,]+, i1 true\)", asm)
@test occursin(r"@air\.simdgroup_matrix_8x8_store\.v64f32\.p1f32\(<64 x float> [^,]+, [^,]+ [^,]+, i64 %[^,]+, <2 x i64> %[^,]+, i1 true\)", asm)
end

# the legacy form is left alone
@test occursin(r"i64 16, <2 x i64> <i64 2, i64 3>, i1 (?:zeroext )?true\)",
air_module(legacy_kernel, v"2.7"))
end

@testset "tensor lowering" begin
# the AIR 2.9 form Metal.jl emits
function kernel(d::Core.LLVMPtr{Int8,1}, p::Core.LLVMPtr{Float32,1},
e::Core.LLVMPtr{Int32,1}, out::Core.LLVMPtr{Int32,1})
ccall("extern air.init_strided_private_tensor.i32.global", llvmcall, Cvoid,
(Core.LLVMPtr{Int8,1}, Int16, Core.LLVMPtr{Float32,1}, Core.LLVMPtr{Int32,1},
Core.LLVMPtr{Int32,1}, Int8, Core.LLVMPtr{Float32,1}, Int32, Int8),
d, 2, p, e, e, 3, reinterpret(Core.LLVMPtr{Float32,1}, C_NULL), 0, -1)
n = ccall("extern air.get_extent_private_tensor.i32", llvmcall, Int32,
(Core.LLVMPtr{Int8,1}, Int16, Int16, Int16), d, 2, 1, 0)
unsafe_store!(out, n)
return
end
# the AIR 2.8 form, as older versions of Metal.jl emitted it for every target
function legacy_kernel(d::Core.LLVMPtr{Int8,1}, p::Core.LLVMPtr{Float32,1},
e::Core.LLVMPtr{Int32,1}, out::Core.LLVMPtr{Int32,1})
ccall("extern air.init_strided_private_tensor.i32.global", llvmcall, Cvoid,
(Core.LLVMPtr{Int8,1}, Int16, Core.LLVMPtr{Float32,1}, Core.LLVMPtr{Int32,1},
Core.LLVMPtr{Int32,1}, Int8),
d, 2, p, e, e, 3)
n = ccall("extern air.get_extent_private_tensor.i32", llvmcall, Int32,
(Core.LLVMPtr{Int8,1}, Int16, Int16), d, 2, 1)
unsafe_store!(out, n)
return
end
function air_module(f, air)
tt = Tuple{Core.LLVMPtr{Int8,1}, Core.LLVMPtr{Float32,1},
Core.LLVMPtr{Int32,1}, Core.LLVMPtr{Int32,1}}
source = methodinstance(typeof(f), tt, Base.get_world_counter())
target = MetalCompilerTarget(; macos=v"26", metal=v"4.0", air)
job = CompilerJob(source, CompilerConfig(target, Metal.CompilerParams(); kernel=true))
sprint(io->GPUCompiler.code_native(io, job; dump_module=true))
end
nullptr = raw"(?:(?:float|i8|\{\}) addrspace\(1\)\*|ptr addrspace\(1\)) null"

@testset "AIR $air" for air in (v"2.8", v"2.9")
asm = air_module(kernel, air)
if air >= v"2.9"
@test occursin(Regex("call void @air\\.init_strided_private_tensor\\.i32\\.global\\(.*, i8 (?:signext )?3, $nullptr, i32 0, i8 (?:signext )?-1\\)\\n"), asm)
@test occursin(r"call i32 @air\.get_extent_private_tensor\.i32\(.*, i16 (?:signext )?1, i16 (?:signext )?0\)\n", asm)
else
@test occursin(r"call void @air\.init_strided_private_tensor\.i32\.global\(.*, i8 (?:signext )?3\)\n", asm)
@test occursin(r"call i32 @air\.get_extent_private_tensor\.i32\(.*, i16 (?:signext )?1\)\n", asm)
end
end

# the AIR 2.8 form is left alone
asm = air_module(legacy_kernel, v"2.8")
@test occursin(r"call void @air\.init_strided_private_tensor\.i32\.global\(.*, i8 (?:signext )?3\)\n", asm)
@test occursin(r"call i32 @air\.get_extent_private_tensor\.i32\(.*, i16 (?:signext )?1\)\n", asm)
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