Large LocalArray eltypes runs into compiler heuristics

[maleadt]

The following works:

function broken_kernel()
    c_frags = LocalArray{Tuple{16}, CUDA.WMMA.Fragment{16, 16, 16, 8, Float16, CUDA.WMMA.Unspecified, CUDA.WMMA.Accumulator}}(undef)

    frag = CUDA.WMMA.Fragment{16, 16, 16, 8, Float16, CUDA.WMMA.Unspecified, CUDA.WMMA.Accumulator}(ntuple(_->Float16(0), 8))
    setindex(c_frags, frag, 1)

    return
end

CUDA.code_llvm(GemmKernels.Kernel.broken_kernel, Tuple{})

But bumping the eltype to Tuple{64} results in an apply_iterate.

[maleadt]

Apparently I fixed that in JuliaLang/julia#46050, so 1.8 is supported. I've added it to CI.

[maleadt]

Let's re-open this to keep track of the max fragment size though. cc @wardvermeulen

[maleadt]

Updated MWE:

using GemmKernels, CUDA

using GemmKernels: LocalArray
using Base: setindex

function kernel()
    c_frags = LocalArray{Tuple{64}, Float32, 1, 64}(undef)
    setindex(c_frags, 0f0, 1)
    return
end

function main()
    CUDA.code_llvm(kernel, Tuple{})
end

isinteractive() || main()

The apply iterate can be avoided by setting the tuple_splat inference and optimization params (max_tuple_splat on 1.10) to a higher value. It defaults to 32. @wardvermeulen Where in the FPU operator implementation did you account for this?

[wardvermeulen]

The apply iterate can be avoided by setting the tuple_splat inference and optimization params (max_tuple_splat on 1.10) to a higher value. It defaults to 32. @wardvermeulen Where in the FPU operator implementation did you account for this?

If you are referring to providing safeguards so this behavior does not occur, I did not account for this in the implementation.

[maleadt]

Ah OK, I thought there were some hard-coded limits that relate to the 16-element LocalArray limit.