Skip to content

Metal: expand vector reductions into scalar chains #962

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

Open
timesselens wants to merge 1 commit into
base: main
Choose a base branch
from
+106 −10
Open

Metal: expand vector reductions into scalar chains #962

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
68 changes: 58 additions & 10 deletions src/metal.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -699,15 +699,6 @@ function lower_air!(@nospecialize(job::CompilerJob{MetalCompilerTarget}), mod::L
end
end

# perform codegen passes that would normally run during machine code emission
if LLVM.has_oldpm()
# XXX: codegen passes don't seem available in the new pass manager yet
@dispose pm=ModulePassManager() begin
expand_reductions!(pm)
run!(pm, mod)
end
end

# flatten chained byte GEPs that the AGX back-end miscompiles for 1-byte accesses on a
# 2-threadgroup grid (see `merge_byte_gep_chains!`). run last, after the intrinsic-lowering
# cleanup above, so the merged form is what reaches the AIR downgrader / back-end.
Expand Down Expand Up @@ -1987,6 +1978,61 @@ function scalarize_vector_minmax!(fun::LLVM.Function)
end
end

# AIR has no vector reductions either, and Apple's back-end fails on `llvm.vector.reduce.*`
# (which LLVM's vectorizers form, e.g. from `prod(size(A))`) or crashes its compiler service.
# LLVM's `ExpandReductions` is a codegen pass that the new pass manager doesn't run, so expand
# each reduction into a chain of its scalar operation over the lanes, in lane order: that is how
# the ordered `fadd`/`fmul` reductions are defined, and a valid order for the others. Min/max
# chain the scalar intrinsic, which the per-call lowering maps to AIR (`and`/`or` on `i1` lanes).
const VECTOR_REDUCTIONS = Dict(
"llvm.vector.reduce.add" => add!,
"llvm.vector.reduce.mul" => mul!,
"llvm.vector.reduce.and" => and!,
"llvm.vector.reduce.or" => or!,
"llvm.vector.reduce.xor" => xor!,
"llvm.vector.reduce.fadd" => fadd!,
"llvm.vector.reduce.fmul" => fmul!,
"llvm.vector.reduce.smax" => "llvm.smax",
"llvm.vector.reduce.smin" => "llvm.smin",
"llvm.vector.reduce.umax" => "llvm.umax",
"llvm.vector.reduce.umin" => "llvm.umin",
"llvm.vector.reduce.fmax" => "llvm.maxnum",
"llvm.vector.reduce.fmin" => "llvm.minnum",
"llvm.vector.reduce.fmaximum" => "llvm.maximum", # LLVM 17+
"llvm.vector.reduce.fminimum" => "llvm.minimum", # LLVM 17+
)
# integer min/max on `i1` lanes, where true is -1 when signed
const I1_MINMAX = Dict("llvm.smax" => and!, "llvm.smin" => or!, "llvm.umax" => or!, "llvm.umin" => and!)

function expand_vector_reductions!(fun::LLVM.Function)
names = filter(n -> LLVM.version() >= v"17" || !endswith(n, "imum"), collect(keys(VECTOR_REDUCTIONS)))
reductions = Dict(LLVM.Intrinsic(n) => VECTOR_REDUCTIONS[n] for n in names)
mod = LLVM.parent(fun)
return lower_intrinsic_calls!(fun) do builder, call, intr
op = get(reductions, intr, nothing)
op === nothing && return nothing
args = collect(LLVM.Value, arguments(call))
vec = last(args) # `fadd`/`fmul` take a start value first
elty = eltype(value_type(vec))
if op isa String && elty == LLVM.Int1Type()
op = I1_MINMAX[op]
elseif op isa String
f = LLVM.Function(mod, LLVM.Intrinsic(op), LLVMType[elty])
op = (builder, a, b) -> call!(builder, function_type(f), f, LLVM.Value[a, b])
end
fmf = elty isa LLVM.FloatingPointType ? LLVM.fast_math(call) : nothing
res = length(args) == 2 ? args[1] : nothing
for i in 0:Int(length(value_type(vec)))-1
lane = extract_element!(builder, vec, ConstantInt(LLVM.Int32Type(), i))
res === nothing && (res = lane; continue)
res = op(builder, res, lane)
# (steps on constants fold to constants, which carry no flags)
fmf !== nothing && res isa LLVM.Instruction && LLVM.fast_math!(res; fmf...)
end
res
end
end

# floating-point math intrinsics that Julia emits as plain `llvm.*` and that Metal exposes as
# AIR device functions. Each has a precise `air.<op>` for f16/f32; some additionally have a
# relaxed, f32-only `air.fast_<op>` that we select when the call is `afn`-flagged — set per-op
Expand Down Expand Up @@ -2396,8 +2442,10 @@ end
function lower_llvm_intrinsics!(@nospecialize(job::CompilerJob), fun::LLVM.Function)
isdeclaration(fun) && return false

# AIR lacks vector min/max intrinsics; scalarize so the per-call lowering below applies.
# AIR lacks vector min/max intrinsics and vector reductions; scalarize them so the per-call
# lowering below applies.
changed = scalarize_vector_minmax!(fun)
changed |= expand_vector_reductions!(fun)

# lower the floating-point math intrinsics Julia emits (sqrt, fma, floor, ...) to their
# AIR device functions, picking the relaxed `air.fast_*` variant for `afn`-flagged calls.
Expand Down
48 changes: 48 additions & 0 deletions test/metal.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1288,6 +1288,54 @@ end
end
end

@testset "vector reduction lowering" begin
# AIR has no vector reductions: each is expanded into a chain of its scalar operation over the
# lanes, in lane order, and min/max into the scalar intrinsics, which are lowered to AIR.
km = @eval module $(gensym())
f() = return
prod4(v) = ccall("llvm.vector.reduce.mul.v4i64", llvmcall, Int64, (NTuple{4, VecElement{Int64}},), v)
end
job, _ = Metal.create_job(km.f, Tuple{})
ops = ["add", "mul", "and", "or", "xor", "smax", "smin", "umax", "umin", "fadd", "fmul", "fmax", "fmin"]
LLVM.version() >= v"17" && append!(ops, ["fmaximum", "fminimum"])
Context() do ctx
ir = IOBuffer()
for op in ops, n in (3, 4)
t, s = op[1] == 'f' ? ("float", "f32") : ("i32", "i32")
args = op in ("fadd", "fmul") ? "$t %s, <$n x $t> %v" : "<$n x $t> %v"
flags = op[1] == 'f' ? "nnan " : ""
println(ir, "declare $t @llvm.vector.reduce.$op.v$n$s($args)")
println(ir, "define $t @$op$n($args) {\n %r = call $flags$t @llvm.vector.reduce.$op.v$n$s($args)\n ret $t %r\n}")
end
println(ir, "declare i1 @llvm.vector.reduce.smax.v3i1(<3 x i1>)")
println(ir, "define i1 @smax_i1(<3 x i1> %v) {\n %r = call i1 @llvm.vector.reduce.smax.v3i1(<3 x i1> %v)\n ret i1 %r\n}")
mod = parse(LLVM.Module, String(take!(ir)))
insts(f) = [i for bb in blocks(f) for i in instructions(bb)]
callees(f) = [LLVM.name(called_operand(i)) for i in insts(f) if i isa LLVM.CallBase]
for f in functions(mod)
isdeclaration(f) && continue
GPUCompiler.lower_llvm_intrinsics!(job, f)
@test !any(startswith("llvm."), callees(f))
end
# the ordered reduction continues from its start value, with the call's fast-math flags
f = functions(mod)["fadd3"]
adds = filter(i -> i isa LLVM.FAddInst, insts(f))
@test length(adds) == 3 && operands(first(adds))[1] == parameters(f)[1]
@test all(i -> LLVM.fast_math(i).nnan, adds)
@test "air.fmin.f32" in callees(functions(mod)["fmin4"])
@test isempty(callees(functions(mod)["smax_i1"])) # `and` on i1 lanes
@test (verify(mod); true)
end

# end to end
@test @filecheck begin
@check_label "define i64 @{{(julia|j)_prod4_[0-9]+}}"
@check_not "@llvm.vector.reduce"
@check_count 3 "mul i64"
Metal.code_native(km.prod4, Tuple{NTuple{4, VecElement{Int64}}})
end
end

@testset "integer intrinsic lowering" begin
# The integer ops Julia emits as llvm.* are lowered to their AIR builtins, so Metal.jl need
# not wrap them. Names/signatures verified against Apple's frontend:
Expand Down
Loading