Smoke: Enzyme

In [1]:

using CUDA, KernelAbstractions, Enzyme
@kernel function square!(y, @Const(x))
    i = @index(Global)
    @inbounds y[i] = x[i]^2
end
f!(y, x) = (square!(get_backend(x))(y, x; ndrange=length(x)); nothing)
x = CuArray(Float32[1, 2, 3, 4]); dx = CUDA.zeros(Float32, 4)
y = CUDA.zeros(Float32, 4);       dy = CUDA.ones(Float32, 4)
Enzyme.autodiff(Reverse, f!, Const, Duplicated(y, dy), Duplicated(x, dx))
Array(dx)

4-element Vector{Float32}:
 2.0
 4.0
 6.0
 8.0