[2.7.0] Metal - llvm.vector.reduce.* reaching the AIR back-end fails to compile

[ArchHem]

Hi, ran into another issue on Metal when using Enzyme autodiff_deferred inside a metal kernel.

After some experimentation, I could reduce it to the following llvmcall example:

using Metal


const Vec4 = NTuple{4,VecElement{Float32}}

hsum(x::NTuple{4,Float32}) =
    ccall("llvm.vector.reduce.fadd.v4f32", llvmcall, Float32, (Float32, Vec4), -0f0, Vec4(map(VecElement, x)))

function cpu(out, xs)
    for i in eachindex(xs)
        @inbounds out[i] = hsum(xs[i])
    end
    return
end

function kernel(out, xs)
    i = Metal.thread_position_in_grid_1d()
    @inbounds out[i] = hsum(xs[i])
    return
end

N = 128
xs = [(1f0, 1f0, 2f0, 3f0) for _ in 1:N]

out = zeros(Float32, N)
cpu(out, xs)
println("CPU ", out[1])

d_xs = MtlArray(xs); d_out = MtlArray(zeros(Float32, N))
try
    Metal.@metal threads=N kernel(d_out, d_xs)
    println("GPU ", Array(d_out)[1])
catch
    println("GPU FAIL")
    for (exc, _) in current_exceptions()
        println("    ", replace(sprint(showerror, exc), "\n" => "\n    "))
    end
end


try
    Metal.@device_code_llvm dump_module=false debuginfo=:none Metal.@metal threads=N kernel(d_out, d_xs)
catch
end
#CPU 7.0
#GPU FAIL
#    NSError: Compiler encountered an internal error (AGXMetalG13X, code 3)
#    Compilation to native code failed; see below for details.
#    If you think this is a bug, please file an issue and attach:
#    - /var/folders/8z/xyt_bsw90n352p3nfj7dfdnc0000gn/T/jl_VMtgz0r272.air
#    - /var/folders/8z/xyt_bsw90n352p3nfj7dfdnc0000gn/T/jl_VMtgz0r272.metallib

With return:

; GPUCompiler.CompilerJob{GPUCompiler.MetalCompilerTarget, Metal.MetalCompilerParams}(MethodInstance for kernel(::MtlDeviceVector{Float32, 1}, ::MtlDeviceVector{NTuple{4, Float32}, 1}), CompilerConfig for GPUCompiler.MetalCompilerTarget, 0x000000000000978f)
[ Info: Load Highlights.jl (`using Highlights`) to enable syntax highlighting of this output.
define void @_Z6kernel14MtlDeviceArrayI7Float32Li1ELi1EES_I5TupleIS0_S0_S0_S0_ELi1ELi1EE(ptr addrspace(1) %state_ptr, ptr addrspace(1) %"out::MtlDeviceArray", ptr addrspace(1) %"xs::MtlDeviceArray", <3 x i32> %thread_position_in_grid) local_unnamed_addr {
conversion:
  %.unpack9 = load ptr addrspace(1), ptr addrspace(1) %"out::MtlDeviceArray", align 8
  %.unpack13 = load ptr addrspace(1), ptr addrspace(1) %"xs::MtlDeviceArray", align 8
  %"[1]" = extractelement <3 x i32> %thread_position_in_grid, i64 0
  %0 = sext i32 %"[1]" to i64
  %1 = getelementptr inbounds [4 x float], ptr addrspace(1) %.unpack13, i64 %0
  %.unpack = load float, ptr addrspace(1) %1, align 4
  %.elt1 = getelementptr inbounds [4 x float], ptr addrspace(1) %.unpack13, i64 %0, i64 1
  %.unpack2 = load float, ptr addrspace(1) %.elt1, align 4
  %.elt3 = getelementptr inbounds [4 x float], ptr addrspace(1) %.unpack13, i64 %0, i64 2
  %.unpack4 = load float, ptr addrspace(1) %.elt3, align 4
  %.elt5 = getelementptr inbounds [4 x float], ptr addrspace(1) %.unpack13, i64 %0, i64 3
  %.unpack6 = load float, ptr addrspace(1) %.elt5, align 4
  %2 = insertelement <4 x float> poison, float %.unpack, i64 0
  %3 = insertelement <4 x float> %2, float %.unpack2, i64 1
  %4 = insertelement <4 x float> %3, float %.unpack4, i64 2
  %5 = insertelement <4 x float> %4, float %.unpack6, i64 3
  %6 = call float @llvm.vector.reduce.fadd.v4f32(float -0.000000e+00, <4 x float> %5)
  %7 = getelementptr inbounds float, ptr addrspace(1) %.unpack9, i64 %0
  store float %6, ptr addrspace(1) %7, align 4
  ret void
}

I am not sure if this issue belongs here or on Enzyme.jl though - aside from Enzyme I think only SIMD.jl would realistically produce these instructions and seems like the codegen is not expanding the vector reductions.

For reference, the Enzyme MWE looks something like:

using Metal, Enzyme, StaticArrays, LinearAlgebra

const V4 = SVector{4,Float32}
const M = SMatrix{4,4,Float32,16}(1, 2, 3, 4,
                                  2, 5, 6, 7,
                                  3, 6, 8, 9,
                                  4, 7, 9, 10)
quad(x::V4) = dot(x, M * x)                  # symmetric M, so the gradient is 2 M x

function kernel(out, xs)
    i = Metal.thread_position_in_grid_1d()
    (dx,), = Enzyme.autodiff_deferred(Reverse, Const(quad), Active, Active(xs[i]))
    @inbounds out[i] = dx
    return
end

(The above does compile if I set Enzyme.API.fast_math!(false) since it cant rassociate)

Versioninfo:

  [dde4c033] Metal v1.11.0
⌃ [90137ffa] StaticArrays v1.9.21

julia> Metal.versioninfo()
macOS 15.7.4, Darwin 24.6.0

Toolchain:
- Julia: 1.12.5
- LLVM: 18.1.7
- Metal: 3.2 (MSL), 2.7 (AIR), 1.2.8 (metallib)

Julia packages:
- Metal.jl: 1.11.0
- GPUArrays: 11.5.14
- GPUCompiler: 2.7.0
- KernelAbstractions: 0.9.42
- ObjectiveC: 6.0.1
- LLVM: 9.13.1
- LLVMDowngrader_jll: 0.10.0+1

Kernel cache:
- binary archives: enabled
  (1 hits, 0 misses this session)

1 device:
- Apple M1 Pro (16 GPU cores, 1.375 MiB allocated; Apple7, Metal3 family)

Let me know if I should attach the .air/.metallib files.