CUDA.jl kernel errors with Julia 1.9.0-rc3

[jgreener64]

I am on main (4c30ed8) and CUDA 4.1.4. I have a GPU kernel where gradients work on Julia 1.8.5 and error on Julia 1.9.0-rc2.

The error is attached as error.txt due to the size limit. If required I can minimise but it depends on other packages so maybe there is a quicker fix.

[wsmoses]

Yeah I think we need the code that produces this.

[jgreener64]

Minimal reproducer below. This is on Julia 1.9.0-rc3, Enzyme main (4a0cb2b), CUDA 4.2.0 and StaticArrays 1.5.24.

It does not error on Julia 1.8.5 and also if you replace T = Float32 with T = Float64, suggesting some issue with Float32 in static arrays.

using Enzyme, CUDA, StaticArrays, LinearAlgebra

function force(coord_i, coord_j)
    ab = coord_j - coord_i
    c = norm(ab) - 1.0
    f = c * normalize(ab)
    return f
end

function kernel!(fs, coords, is, js)
    inter_i = (blockIdx().x - 1) * blockDim().x + threadIdx().x
    @inbounds if inter_i <= length(is)
        i, j = is[inter_i], js[inter_i]
        f = force(coords[i], coords[j])
        fs[1, i] += f[1]
    end
    return nothing
end

function grad_kernel!(fs, d_fs, coords, d_coords, is, js)
    autodiff_deferred(
        Reverse,
        kernel!,
        Const,
        Duplicated(fs, d_fs),
        Duplicated(coords, d_coords),
        Const(is),
        Const(js),
    )
    return nothing
end

T = Float32
coords = CuArray(rand(SVector{3, T}, 10))
fs = CuArray(zeros(3, 10))
is = CuArray([Int32(i) for i in 1:9])
js = CuArray([Int32(i) for i in 2:10])
d_coords = zero(coords)
d_fs = zero(fs)

CUDA.@sync @cuda threads=10 blocks=1 grad_kernel!(fs, d_fs, coords, d_coords, is, js)

Error in error.txt.

[wsmoses]

@vchuravy this feels like a version of #683 per the sqrtf. However we should hypothetically have resolved this already in preservenvvm from the llvm level.

[jgreener64]

Also an issue with Julia 1.9.0, Enzyme main (96ddde3), CUDA 4.2.0 and StaticArrays 1.5.25. Error in error.txt.

[vchuravy]

@wsmoses tells me this might be fixed by EnzymeAD/Enzyme#1233