Metal: expand vector reductions into scalar chains

[timesselens]

AIR has no vector reductions, and from Julia 1.12 (LLVM 17 and later) nothing expands llvm.vector.reduce.* before Apple's back-end, which then fails or crashes its compiler service: lower_air! runs ExpandReductions only with the legacy pass manager. This PR proposes expanding each reduction in lower_llvm_intrinsics!, on every LLVM version.

Before / after

using Metal
prod4(v) = ccall("llvm.vector.reduce.mul.v4i64", llvmcall, Int64, (NTuple{4,VecElement{Int64}},), v)
function k!(out, x)
    i = Metal.thread_position_in_grid_1d()
    @inbounds out[i] = prod4(ntuple(j -> VecElement(x[4(i-1)+j]), Val(4)))
    return
end
x = MtlArray(Int64.(1:16)); out = MtlArray{Int64}(undef, 4)
@metal threads=4 k!(out, x); Array(out)
# main (d36ab74, Metal.jl 1.11.1, Julia 1.12.7): ERROR: Compilation to native code failed
# this PR:                                       [24, 1680, 11880, 43680]

The intrinsic came up in an Enzyme reverse-mode kernel. Enzyme.jl ran LLVM's vectorizers over the differentiated module, which GPUCompiler doesn't do for Metal, and the SLP vectorizer turned the length of a 4-D array into llvm.vector.reduce.mul.v4i64; that kernel crashed Apple's compiler service (XPC_ERROR_CONNECTION_INTERRUPTED). EnzymeAD/Enzyme.jl#3659 stops that at the source. This PR still covers reductions from other sources, as #549 did.

Proposed change

expand_vector_reductions!, next to scalarize_vector_minmax!, replaces each reduction with a chain of its scalar operation over the lanes, in lane order: that is how the ordered fadd/fmul reductions are defined (from their start value), and a valid order for the others. Min/max chain the scalar intrinsic, which the per-call lowering then maps to AIR; on i1 lanes they become and/or. The reduction's fast-math flags carry over. The legacy expand_reductions! call goes, so every LLVM version takes the same path. LLVM's own pass can't take over there: the new pass manager doesn't know it (unknown function pass 'expand-reductions' on LLVM 18 and 20).

Tests

"vector reduction lowering" in test/metal.jl, at the IR level: every reduction on 3 and 4 lanes (and smax on i1 lanes) leaves no llvm.* call after lower_llvm_intrinsics!; the ordered fadd continues from its start value with the call's fast-math flags; fmin becomes air.fmin.f32; and code_native of llvm.vector.reduce.mul.v4i64 gives three mul i64. On main the new set fails.

Open questions

• Reductions over i1 lanes that Julia's optimizer turns into bitcast <4 x i1> to i4 + icmp (e.g. any over <4 x i1>) still fail, with or without this PR, because Apple's back-end rejects the i4. Is odd-width integer legalization already on your list?

Related

• Metal: Run ExpandReductions pass. #549 added the ExpandReductions call this PR replaces.
• macOS 27: bare LLVM fence crashes the AGX back-end compiler (XPC_ERROR_CONNECTION_INTERRUPTED) Metal.jl#968: the same compiler-service crash, from bare fences.

Verification

• Pkg.test(test_args=["metal"]) on macOS aarch64: Julia 1.10.12 and 1.11.9 (also with --opaque-pointers), 1.12.7 and 1.13.0 all pass.
• On an M1 Max (Metal.jl 1.11.1), 45 kernels, one per reduction and lane type (Int64×4, Int32×3, Int16×5; Float32×4, Float32×5, Float16×3), compile, run, and match the host's evaluation of the same intrinsic.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 86.75%. Comparing base (d36ab74) to head (db7185b).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #962      +/-   ##
==========================================
+ Coverage   86.65%   86.75%   +0.09%     
==========================================
  Files          29       29              
  Lines        5786     5807      +21     
==========================================
+ Hits         5014     5038      +24     
+ Misses        772      769       -3     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.