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