# KernelAbstractions + Enzyme - how to do GPU-side autodiff?

**URL:** <https://discourse.julialang.org/t/kernelabstractions-enzyme-how-to-do-gpu-side-autodiff/132643>\
**Category:** GPU\
**Tags:** question, autodiff, enzyme\
**Created:** [September 24, 2025, 11:07pm UTC](https://discourse.julialang.org/t/kernelabstractions-enzyme-how-to-do-gpu-side-autodiff/132643 "2025-09-24T23:07:27Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![jl\_enthusiast](https://avatars.discourse-cdn.com/v4/letter/j/7993a0/32.png) [@jl\_enthusiast](https://discourse.julialang.org/u/jl_enthusiast)\
**Post date:** [September 24, 2025, 11:07pm UTC](https://discourse.julialang.org/t/kernelabstractions-enzyme-how-to-do-gpu-side-autodiff/132643/1 "2025-09-24T23:07:27Z")

</div>

Hello,

I recently tried to (and failed) to auto-differentiate kernels with Enzyme.jl, written using KernelAbstractions.jl.

I tried to search for examples in other packages, but could not find any concrete examples of Enzyme + KA being used together (the Enzyme test cases are written in CUDA/specific backends [Enzyme.jl/test at main · EnzymeAD/Enzyme.jl · GitHub](https://github.com/EnzymeAD/Enzyme.jl/tree/main/test)) and as such have no idea how it might look for KA - I assume I would need to launch a separate kernel, with a deferred autodiff call inside?

My current (very simple) attempt using KA was

```julia-auto
using KernelAbstractions
using Enzyme

@kernel function cos_kernel_ka!(B, A)
    i = @index(Global, Linear)
    B[i] = sin(A[i])
end
@kernel function grad_cos_kernel_ka!(dA, dB, A, B)
    Enzyme.autodiff_deferred(
        Reverse,
        Const(cos_kernel_ka!),
        Const,
        Duplicated(B, dB),
        Duplicated(A, dA)
    )
end

```

Which results in a `ERROR: AssertionError: actualRetType != Union{}` on Metal/CUDA backends when I attempt to launch it, so I strongly suspect I just have a syntax error.

Appreciate any help in advance!

---

<div class="post-metadata">

**Author:** ![yolhan\_mannes](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/yolhan_mannes/32/220485_2.png) [@yolhan\_mannes](https://discourse.julialang.org/u/yolhan_mannes)\
**Post date:** [September 25, 2025, 6:34am UTC](https://discourse.julialang.org/t/kernelabstractions-enzyme-how-to-do-gpu-side-autodiff/132643/2 "2025-09-25T06:34:49Z")

</div>

Why you make your gradient a kernel ? Kernel only supports standard operations for now.  
In your case the autodiff will make a kernel in a kernel which isn’t supported right now.  
Also thats not how the kernel is run. However I don’t what you want to calculate :

```julia-auto
julia> using KernelAbstractions,Enzyme,CUDA
julia> @kernel function cos_kernel_ka!(B, A)
           i = @index(Global, Linear)
           B[i] = sin(A[i])
       end

julia> function cos_run!(B,A)
           cos_kernel_ka!(get_backend(B))(B,A;ndrange=length(B))
           B
       end
cos_run! (generic function with 1 method)

julia> function grad_cos_kernel_ka!(dA, dB, A, B)
           Enzyme.autodiff_deferred(
               Reverse,
               Const(cos_run!),
               Active,
               DuplicatedNoNeed(B, dB),
               Duplicated(A, dA)
           )
       end
grad_cos_kernel_ka! (generic function with 1 method)

julia> A = CUDA.rand(100);

julia> B = CUDA.rand(100);

julia> dA = Enzyme.make_zero(A);

julia> dB = Enzyme.make_zero(B);

julia> grad_cos_kernel_ka!(dA,dB,A,B)
ERROR: Return type of differentiated function was not a scalar as required, found Float32[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 … 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
If calling Enzyme.autodiff(Reverse, f, Active, ...), try Enzyme.autodiff_thunk(Reverse, f, Duplicated, ....)
If calling Enzyme.gradient, try Enzyme.jacobian

```

this works :

```julia-auto
julia> using KernelAbstractions,Enzyme,CUDA

julia> A = CUDA.rand(100);

julia> B = CUDA.rand(100);

julia> dA = Enzyme.make_zero(A);

julia> dB = Enzyme.make_zero(B);

julia> @kernel function cos_kernel_ka!(B, A)
           i = @index(Global, Linear)
           B[i] = sin(A[i])
       end
cos_kernel_ka! (generic function with 4 methods)

julia> function cos_run!(B,A)
           cos_kernel_ka!(get_backend(B))(B,A;ndrange=length(B))
           sum(B)
       end
cos_run! (generic function with 1 method)

julia> function grad_cos_kernel_ka!(dA, dB, A, B)
          Enzyme.autodiff_deferred(
               Reverse,
               Const(cos_run!),
               Active,
               DuplicatedNoNeed(B, dB),
               Duplicated(A, dA)
           )
       end
grad_cos_kernel_ka! (generic function with 1 method)

julia> grad_cos_kernel_ka!(dA,dB,A,B)
((nothing, nothing),)

julia> dA
100-element CuArray{Float32, 1, CUDA.DeviceMemory}:
 0.573843
 0.68915915
 0.9999656
 0.9991022
 0.92152673
 0.5781501
 0.97622705
 0.9257504
 0.94532764
 0.85892415
 0.7805012
 0.7845093
 0.9911921
 ⋮
 0.8982199
 0.98546714
 0.91434866
 0.9942696
 0.8421704
 0.71289515
 0.9669826
 0.9899486
 0.9999464
 0.9949909
 0.7747586
 0.98301524

```

which is normal actually what did you want to calculate ? If you want the jacobian or the v^T J product, reffer to Enzyme docs or the implementation in jacobian or just use jacobian itself

ps : depending on what you want don’t forget to zeros dA and dB

```julia-auto
function grad_cos_kernel_ka!(dA, dB, A, B)
           Enzyme.make_zero!(dA)
           Enzyme.make_zero!(dB)
           Enzyme.autodiff_deferred(
               Reverse,
               Const(cos_run!),
               Active,
               DuplicatedNoNeed(B, dB),
               Duplicated(A, dA)
           )
       end

```

ps 2 : you can do this

```julia-auto
julia> using KernelAbstractions,Enzyme,CUDA

julia> A = CUDA.rand(100);

julia> B = CUDA.rand(100);

julia> dA = Enzyme.make_zero(A);

julia> dB = Enzyme.make_zero(B);

julia> @kernel function cos_kernel_ka!(B, A)
           i = @index(Global, Linear)
           B[i] = sin(A[i])
       end

julia> function cos_run!(B,A)
           cos_kernel_ka!(get_backend(B))(B,A;ndrange=length(B))
           nothing
       end
cos_run! (generic function with 1 method)

julia> function grad_cos_kernel_ka!(dA, dB, A, B)
           Enzyme.autodiff_deferred(
               Reverse,
               Const(cos_run!),
               Const,
               DuplicatedNoNeed(B, dB),
               Duplicated(A, dA)
           )
       end
grad_cos_kernel_ka! (generic function with 1 method)

julia> dB .= 1;

julia> grad_cos_kernel_ka!(dA,dB,A,B)
((nothing, nothing),)

julia> dA ≈ cos.(A)
true

```