JuliaGPU / JuliaGPU/KernelAbstractions.jl

issue with using enzyme on kernel abstractions on CUDA backend

Open
#480 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
523
Forks
88
Avg merge
1d 11h
Merged PRs (30d)
25

Description

I want to desing rrule from chainrules for my kernel. Below is simple reproducible example.
System info:
Julia 1.10
CUDA v5.4.0
ChainRulesCore v1.23.0
ChainRulesTestUtils v1.13.0
Enzyme v0.12.9 https://github.com/EnzymeAD/Enzyme.jl.git#main
EnzymeTestUtils v0.1.7
KernelAbstractions v0.9.20 https://github.com/JuliaGPU/KernelAbstractions.jl#main

GPU: Nvidia RTX 3090

code

using KernelAbstractions
using ChainRulesCore, Zygote, CUDA, Enzyme, Test

@kernel function example_kenr(@Const(A),A_out)

    index = @index(Global)
    shared_arr = @localmem Float32 (@groupsize()[1], 1)
    shared_arr[@index(Local, Linear)] = A[index]
    A_out[index] = shared_arr[@index(Local, Linear), 1]
    index = @index(Global)
end

function call_example(A,A_out)
    dev = get_backend(A)
    example_kenr(dev, 256)(A,A_out, ndrange=(size(A)[1]))
    KernelAbstractions.synchronize(dev)
    return nothing
end


A=CUDA.ones(10).*2
A_out=CUDA.ones(10)
call_example(A,A_out)
@test A_out == CUDA.ones(10).*2


function ChainRulesCore.rrule(::typeof(call_example), A,A_out)

    #modify A_out by mutation
    call_example(A,A_out)

    function call_test_kernel1_pullback(d_A_out)
        d_A_out = CuArray(collect(d_A_out))
        d_A = CUDA.zeros(size(A)...)

        Enzyme.autodiff_deferred(Enzyme.Reverse, call_example, Const, Duplicated(A,d_A), Duplicated(A_out, d_A_out))

        #NoTangent for the function itself
        return NoTangent(), d_A,d_A_out
    end


    return A_out, call_test_kernel1_pullback

end

out,pull_back=rrule(call_example,A,A_out)
pull_back(CUDA.ones(10))

error

ERROR: Enzyme.Compiler.EnzymeRuntimeException(Cstring(0x00007d0b6b39b32c))
Stacktrace:
  [1] throwerr(cstr::Cstring)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/sDjFs/src/compiler.jl:1338
  [2] getindex
    @ ./essentials.jl:13 [inlined]
  [3] get
    @ ./dict.jl:525 [inlined]
  [4] compiler_cache
    @ ~/.julia/packages/CUDA/DS19C/src/compiler/compilation.jl:166 [inlined]
  [5] macro expansion
    @ ~/.julia/packages/CUDA/DS19C/src/compiler/execution.jl:366 [inlined]
  [6] macro expansion
    @ ./lock.jl:267 [inlined]
  [7] #cufunction#1169
    @ ~/.julia/packages/CUDA/DS19C/src/compiler/execution.jl:364
  [8] cufunction
    @ ~/.julia/packages/CUDA/DS19C/src/compiler/execution.jl:361 [inlined]
  [9] macro expansion
    @ ~/.julia/packages/CUDA/DS19C/src/compiler/execution.jl:112 [inlined]
 [10] #_#4
    @ ~/.julia/packages/CUDA/DS19C/src/CUDAKernels.jl:103 [inlined]
 [11] augmented_julia____4_10560_inner_1wrap
    @ ~/.julia/packages/CUDA/DS19C/src/CUDAKernels.jl:0
 [12] macro expansion
    @ ~/.julia/packages/Enzyme/sDjFs/src/compiler.jl:5916 [inlined]
 [13] enzyme_call
    @ ~/.julia/packages/Enzyme/sDjFs/src/compiler.jl:5566 [inlined]
 [14] AugmentedForwardThunk
    @ ~/.julia/packages/Enzyme/sDjFs/src/compiler.jl:5454 [inlined]
 [15] runtime_generic_augfwd(activity::Type{…}, width::Val{…}, ModifiedBetween::Val{…}, RT::Val{…}, f::CUDA.CUDAKernels.var"##_#4", df::Nothing, primal_1::Int64, shadow_1_1::Nothing, primal_2::Nothing, shadow_2_1::Nothing, primal_3::KernelAbstractions.Kernel{…}, shadow_3_1::Nothing, primal_4::CuArray{…}, shadow_4_1::CuArray{…}, primal_5::CuArray{…}, shadow_5_1::CuArray{…})
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/sDjFs/src/rules/jitrules.jl:179
 [16] Kernel
    @ ~/.julia/packages/CUDA/DS19C/src/CUDAKernels.jl:89 [inlined]
 [17] call_example
    @ ~/projects_new/superVoxelJuliaCode/superVoxelJuliaCode/src/old/cuda_enzyme_kern_ans_test.jl:15 [inlined]
 [18] diffejulia_call_example_4822wrap
    @ ~/projects_new/superVoxelJuliaCode/superVoxelJuliaCode/src/old/cuda_enzyme_kern_ans_test.jl:0
 [19] macro expansion
    @ ~/.julia/packages/Enzyme/sDjFs/src/compiler.jl:5916 [inlined]
 [20] enzyme_call
    @ ~/.julia/packages/Enzyme/sDjFs/src/compiler.jl:5566 [inlined]
 [21] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/sDjFs/src/compiler.jl:5443 [inlined]
 [22] autodiff_deferred
    @ ~/.julia/packages/Enzyme/sDjFs/src/Enzyme.jl:440 [inlined]
 [23] autodiff_deferred
    @ ~/.julia/packages/Enzyme/sDjFs/src/Enzyme.jl:510 [inlined]
 [24] (::var"#call_test_kernel1_pullback#5"{CuArray{…}, CuArray{…}})(d_A_out::CuArray{Float32, 1, CUDA.DeviceMemory})
    @ Main ~/projects_new/superVoxelJuliaCode/superVoxelJuliaCode/src/old/cuda_enzyme_kern_ans_test.jl:36
 [25] top-level scope
    @ ~/projects_new/superVoxelJuliaCode/superVoxelJuliaCode/src/old/cuda_enzyme_kern_ans_test.jl:48
Some type information was truncated. Use `show(err)` to see complete types.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the CUDA, KernelAbstractions, ChainRulesCore, and Enzyme reproducer in src/old/cuda_enzyme_kern_ans_test.jl, then inspect the stack through CUDAKernels.cufunction and Enzyme.autodiff_deferred. Done means the pullback completes without the Enzyme.Compiler.EnzymeRuntimeException and the existing A_out test still passes.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
hpc
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.