[KA + Enzyme]: Instruction does not dominate all uses!

[michel2323]

The following code crashes when differentiating through a KA CPU() kernel. This code will eventually be the unit test code. Without the struct, everything works fine.

A reproducer with the right branches in the Manifest.toml can be found here (Edit: Not needed anymore. See https://github.com/EnzymeAD/Enzyme.jl/issues/1402#issuecomment-2113466008) log.txt

using Test
using Enzyme
using KernelAbstractions
using EnzymeCore
using EnzymeCore.EnzymeRules

const KA = KernelAbstractions

struct MyData
    A::Array{Float64}
    B::Array{Float64}
end

@kernel function square!(A,B)
    I = @index(Global, Linear)
    @inbounds A[I] = B[I] * B[I]
end

function square_caller(data, backend)
    kernel = square!(backend)
    kernel(data.A, data.B, ndrange=size(data.A))
    return nothing
end

function square_caller(A, B, backend)
    kernel = square!(backend)
    kernel(A, B, ndrange=size(A))
    return nothing
end

function enzyme_testsuite(backend, ArrayT)
    A = ArrayT(zeros(64))
    dA = ArrayT{Float64}(undef, 64)
    B = ArrayT{Float64}(undef, 64)
    dB = ArrayT{Float64}(undef, 64)

    dA .= 1
    B .= (1:1:64)
    dB .= 1
    data = MyData(A, B)
    ddata = MyData(dA, dB)
    square_caller(data, backend())
    Enzyme.autodiff(
        Reverse, square_caller, Duplicated(A, dA), Duplicated(B, dB), Const(backend())
    )
    KA.synchronize(backend())
    # Does not work
    Enzyme.autodiff(
        Reverse, square_caller, Duplicated(data, ddata), Const(backend())
    )
    KA.synchronize(backend())
    # @show ddata.B
    # @show dB
end

enzyme_testsuite(CPU, Array)
# enzyme_testsuite(CUDABackend, CuArray)

[wsmoses]

@michel2323 can you check that this still errs?

I believe the root cause of the issue should be gone now.

[michel2323]

@wsmoses This still breaks the same way on KA#main and Enzyme#main (self-compiled [email protected]). No CUDA needed. Removed the synchronize in the kernel to not be dependent on the KA GPU PR.

[wsmoses]

@michel2323 this works for me on Enzyme#main. Can you confirm it still fails (and post a current full error log and all relevant julia/package/os versions)