Enzyme gradient always seems to stall

Whatever we try, me and @VasylHafych always run into a stall when trying to autodiff kernels with Enzyme.

The following stalls at evt = dmatmul(...) (example taken from https://github.com/JuliaGPU/KernelAbstractions.jl/blob/master/lib/KernelGradients/test/matmul.jl):

using KernelAbstractions, Enzyme, Test


# From KernelGradients.jl:
function Enzyme.autodiff(kernel::KernelAbstractions.Kernel{<:Any, <:Any, <:Any, Fun}) where Fun
    f = kernel.f
    function df(ctx, args...)
        Enzyme.autodiff_deferred(f::Fun, Enzyme.Const, ctx, args...)
    end
    similar(kernel, df)
end


@kernel function matmul_kernel!(a, b, c)
    i, j = @index(Global, NTuple)

    # creating a temporary sum variable for matrix multiplication
    tmp_sum = zero(eltype(c))
    for k = 1:size(a)[2]
        @inbounds tmp_sum += a[i,k] * b[k, j]
    end

    c[i,j] = tmp_sum
end


ArrayT = Array
a = ArrayT(rand(128, 256))
b = ArrayT(rand(256, 128))
c = ArrayT(zeros(128, 128))
dev = CPU()

matmul = matmul_kernel!(dev, (32, 32))
wait(matmul(a, b, c, ndrange=size(c)))

@test c ≈ a*b

dmatmul = Enzyme.autodiff(matmul)
da = similar(a)
da .= 0
db = similar(b)
db .= 0
dc = similar(c)
dc .= 1
c .= 0

compare_dc = copy(dc)

evt = dmatmul(Duplicated(a, da), Duplicated(b, db), Duplicated(c, dc), ndrange=size(c))

# wait(evt)
# @test da ≈ compare_dc * b'
# @test db ≈ a' * compare_dc

We've tried KernelAbstractions v0.8.0 with the current Enzyme main branch, as well as KA v0.7 with KernelGradients and the last Enzyme release, etc.

Reproduce the Julia example in the issue, starting with lib/KernelGradients/test/matmul.jl and the Enzyme.autodiff integration shown there. Investigate why execution stalls at dmatmul for the listed KernelAbstractions and Enzyme versions. Done means dmatmul returns and the commented gradient checks for da and db can run successfully.