From c60a69df5e7a24e7703fcff5a80cbcd778782931 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 21 Oct 2026 07:00:47 +0200 Subject: [PATCH] Metal: legalize the simdgroup matrix and tensor intrinsics Metal.jl emits these in their newest form and rewrote them for older AIR versions in its own `finish_ir!`. Like the barrier and atomic ABIs, that is a property of the AIR version being targeted, so do it here during AIR lowering: before AIR 2.8, the simdgroup matrix load/store calls get the legacy elements-per-row, origin and transpose operands, and before AIR 2.9 the operands AIR 2.9 appended to the tensor intrinsics are dropped. Calls already in the older form, as emitted by released versions of Metal.jl, are left alone. The element type of the simdgroup matrix pointer operand is now also derived here, for the downgrader. --- src/metal.jl | 136 ++++++++++++++++++++++++++++++++++++++++++++++++++ test/metal.jl | 126 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 262 insertions(+) diff --git a/src/metal.jl b/src/metal.jl index a7b63cad..0742ddf0 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -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 @@ -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__`, 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) diff --git a/test/metal.jl b/test/metal.jl index 2ed6b034..f1b1d7f4 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -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> , <2 x i64> , <2 x i64> \\)"), asm) + @test occursin(Regex("call void @air\\.simdgroup_matrix_8x8_store\\.v64f32\\.p1f32\\($v64 [^,]+, $ptr [^,]+, <2 x i64> , <2 x i64> , <2 x i64> \\)"), 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> , i1 true\\)"), asm) + @test occursin(Regex("call void @air\\.simdgroup_matrix_8x8_store\\.v64f32\\.p1f32\\($v64 [^,]+, $ptr [^,]+, i64 16, <2 x i64> , 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> , 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