From db7185b8a3862c37da1f3b36f5bcb7b1071fb75d Mon Sep 17 00:00:00 2001 From: Tim Esselens Date: Fri, 25 Sep 2026 14:27:57 +0200 Subject: [PATCH] Metal: expand vector reductions into scalar chains AIR has no vector reductions, 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. `ExpandReductions` only ran with the legacy pass manager, so from LLVM 17 on (Julia 1.12) nothing expanded them. 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). This replaces the legacy `expand_reductions!` call, so every LLVM version takes the same path. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/metal.jl | 68 +++++++++++++++++++++++++++++++++++++++++++-------- test/metal.jl | 48 ++++++++++++++++++++++++++++++++++++ 2 files changed, 106 insertions(+), 10 deletions(-) diff --git a/src/metal.jl b/src/metal.jl index ce4ae7b0..2a83412a 100644 --- a/src/metal.jl +++ b/src/metal.jl @@ -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. @@ -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.` for f16/f32; some additionally have a # relaxed, f32-only `air.fast_` that we select when the call is `afn`-flagged — set per-op @@ -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. diff --git a/test/metal.jl b/test/metal.jl index 9b69b434..a37b4bc3 100644 --- a/test/metal.jl +++ b/test/metal.jl @@ -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: