Skip to content

Leave the AIR legalization of simdgroup and tensor intrinsics to GPUCompiler #1096

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 2 commits into from
+6 −165
Closed

Leave the AIR legalization of simdgroup and tensor intrinsics to GPUCompiler #1096

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
23ef20a
Leave the AIR legalization of the simdgroup and tensor intrinsics to …
maleadt Oct 21, 2026
a943c02
[sources] GPUCompiler.jl tb/metal-intrinsic-abi
maleadt Oct 21, 2026
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
3 changes: 3 additions & 0 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -78,3 +78,6 @@ Statistics = "1"
UUIDs = "1"
UnsafeAtomics = "0.4"
julia = "1.10"

[sources]
GPUCompiler = {url = "https://github.com/JuliaGPU/GPUCompiler.jl", rev = "tb/metal-intrinsic-abi"}
162 changes: 0 additions & 162 deletions src/compiler/compilation.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -164,25 +164,6 @@ function GPUCompiler.finish_module!(@nospecialize(job::MetalCompilerJob),
return entry
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 the simdgroup intrinsics build their dims/strides 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

function is_tensor_op_descriptor_constant(gv::LLVM.GlobalVariable)
# Shader Validation faults if tensor-op descriptors are copied out of AIR's
# constant address space, so leave just those descriptor globals in AS0.
Expand Down Expand Up @@ -235,149 +216,6 @@ function GPUCompiler.finish_ir!(@nospecialize(job::MetalCompilerJob),
Tuple{CompilerJob{MetalCompilerTarget}, LLVM.Module, LLVM.Function},
job, mod, entry)

# downgrade intrinsics when targeting older AIR versions
# (GPUCompiler legalizes atomics, see `lower_atomics!`)
## simdgroup
if job.config.target.air < v"2.8"
# AIR 2.8 generalized the simdgroup matrix load/store intrinsics, replacing
# 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 those vectors' elements (see
# `simdgroup_load/store` in src/device/intrinsics/simd.jl). recover the legacy
# operands from the 2.8 layout: the elements-per-row is the non-unit stride, the
# transpose flag is statically recovered from which stride is unit (the legacy
# intrinsic needs it as an immediate), and the transposed layout's swapped origin
# is unswapped back.
for f in collect(mod.functions)
fn = f.name
m = match(r"^air\.simdgroup_matrix_8x8_(load|store)\.", fn)
m === nothing && continue
is_load = m.captures[1] == "load"

calls = collect(f.users)
@assert all(call -> call isa LLVM.CallInst, calls)

# construct the legacy function type
T_i64 = LLVM.Int64Type()
T_vec2 = LLVM.VectorType(T_i64, 2)
T_bool = LLVM.Int1Type()
new_ft = f.function_type
old_params = if is_load
# value pointer; elements per row; origin; transpose
[new_ft.parameters[1], T_i64, T_vec2, T_bool]
else
# value; value pointer; elements per row; origin; transpose
[new_ft.parameters[1:2]..., T_i64, T_vec2, T_bool]
end
old_ft = LLVM.FunctionType(new_ft.return_type, old_params)

# redeclare the function and rewrite its calls
f.name = fn * ".air28"
old_f = LLVM.Function(mod, fn, old_ft)
# carry over the attributes (convergent etc.); they were attached before
# optimization, so GPUCompiler won't re-derive them for this declaration
append!(old_f.function_attributes, f.function_attributes)
for call in calls
@dispose builder=IRBuilder() begin
# (positioning the builder gives it the call's debug location)
position!(builder, LLVM.before(call))

args = collect(call.arguments)
prefix = is_load ? args[1:1] : args[1:2]
_, strides, origin = args[end-2:end]

# the transposed layout has a unit stride in the second dimension;
# the non-transposed one in the first. one of the two is always the
# literal 1, so this is statically decidable.
transposed = static_vector_lane(strides, 1) == 1
@assert transposed || static_vector_lane(strides, 0) == 1 """
unexpected simdgroup matrix strides $(strides); the AIR downgrade \
only handles the transposed and non-transposed column-major layouts \
emitted by `simdgroup_load`/`simdgroup_store`"""

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

# the transposed layout emits a row/column-swapped origin; unswap it
# for the legacy form, which leaves the non-transposed origin as-is
legacy_origin = origin
if transposed
o1 = extract_element!(builder, origin, ConstantInt(Int32(0)))
o2 = extract_element!(builder, origin, ConstantInt(Int32(1)))
legacy_origin = insert_element!(builder, UndefValue(T_vec2), o2,
ConstantInt(Int32(0)))
legacy_origin = insert_element!(builder, legacy_origin, o1,
ConstantInt(Int32(1)))
end

new_call = call!(builder, old_ft, old_f,
[prefix..., epr, legacy_origin,
ConstantInt(T_bool, transposed)])
replace_uses!(call, new_call)
erase!(call)
end
end

@assert isempty(f.uses)
erase!(f)
end
end
## tensor
if job.config.target.air < v"2.9"
# AIR 2.9 appended operands to these, which the front-end always emits; drop them
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
new_ft = f.function_type
old_ft = LLVM.FunctionType(new_ft.return_type, new_ft.parameters[1:nparams])

f.name = fn * ".air29"
old_f = LLVM.Function(mod, fn, old_ft)
append!(old_f.function_attributes, f.function_attributes)
for call in collect(f.users)
call::LLVM.CallInst
@dispose builder=IRBuilder() begin
position!(builder, LLVM.before(call))
new_call = call!(builder, old_ft, old_f,
collect(call.arguments)[1:nparams])
replace_uses!(call, new_call)
erase!(call)
end
end
erase!(f)
end
end

# pointer type information for typed intrinsics
# (this is consumed by the LLVM IR downgrader)
for (jltyp, llvmtyp) in (Int32 => :i32, Int64 => :i64,
Float16 => :f16, Float32 => :f32,
BFloat16 => :bf16),
(as, asname) in (AS.Device => "global", AS.ThreadGroup => "local")

# map of intrinsics to pointer operand indices and eltypes
intrinsics = Dict()
## simd
intrinsics["simdgroup_matrix_8x8_load.v64$llvmtyp.p$as$llvmtyp"] = (1 => jltyp,)
intrinsics["simdgroup_matrix_8x8_store.v64$llvmtyp.p$as$llvmtyp"] = (2 => jltyp,)

# apply metadata to the function declarations
for (intr, args) in intrinsics
fn = "air.$intr"
f = get(mod.functions, fn, nothing)
f === nothing && continue
mds = []
for (idx, typ) in args
push!(mds, ConstantInt(Int32(idx-1)))
push!(mds, null(convert(LLVMType, typ)))
end
f.metadata["arg_eltypes"] = MDNode(mds)
end
end

# name the kernel readably (see `kernel_symbol`). this has to happen after GPUCompiler
# has looked up the entry point by its original name, i.e., not before optimization.
if job.config.kernel && job.config.name === nothing
Expand Down
4 changes: 2 additions & 2 deletions src/device/intrinsics/simd.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,8 @@ end
# the transposed layout of a column-major matrix is dims = (8, epr), strides = (epr, 1),
# and a row/column-swapped origin, while the non-transposed layout swaps each to
# dims = (epr, 8), strides = (1, epr) and an unswapped origin. when targeting older AIR
# versions, `finish_ir!` rewrites these calls to the legacy elements-per-row + transpose
# flag signature (keep in sync with the downgrade rule in src/compiler/compilation.jl).
# versions, GPUCompiler rewrites these calls to the legacy elements-per-row + transpose
# flag signature, which requires one of the strides to be the literal 1.
for (jltype, suffix) in ((:Float16, "f16"), (:Float32, "f32"), (:BFloat16, "bf16"))
for as in (AS.Device, AS.ThreadGroup)
@eval begin
Expand Down
2 changes: 1 addition & 1 deletion src/device/intrinsics/tensor.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ tensor_datatype(::Type{UInt32}) = Int8(8)
# Build an `i32`-indexed strided tensor view over a device-memory buffer. These are emitted
# in the form of AIR 2.9, which appends a pointer into the same address space as the data,
# an index and an `i8`, for which Apple's frontend always passes `null, 0, -1` (presumably
# describing blockwise-quantization scales, which these views don't have); `finish_ir!`
# describing blockwise-quantization scales, which these views don't have); GPUCompiler
# drops them for AIR 2.8.
@device_function @inline init_strided_tensor_device!(
handle::TensorDescriptorStorage,
Expand Down