JuliaGPU / JuliaGPU/KernelAbstractions.jl

Improve error on Enzyme inability to differentiate

Open
#629 5 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

When trying to run Enzyme autodiff on a kernel, even with the simplest MWE I get the same error. The error doesn't happen when I use CUDA kernels directly. This is the code:

Main code:

using KernelAbstractions
using CUDA
using DifferentiationInterface
using Enzyme
params = (
    γ = Float32(42.58e6 * 2π),
    B = Float32(1e-6),      
    T1 = 1.0f0,     
    T2 = 0.5f0,     
    M0 = 1.0f0,)
@inline function cross2(a, b)
    return (a[2] * b[3] - a[3] * b[2],
            a[3] * b[1] - a[1] * b[3],
            a[1] * b[2] - a[2] * b[1])
end

@kernel function bloch!(dm, m, @Const(params))
    γ, B, T1, T2, M0 = params
    i = @index(Global)
    m_i = @views(m[:, i])
    dm_i = @views(dm[:, i])
    B = (0.0f0, 0.0f0, B) 
    # cross_term = γ .* cross2(m_i, B)
    relax_term = (m_i[1] / T2, m_i[2] / T2, (m_i[3] - M0) / T1)
    # dM = cross_term .- relax_term
    dM = relax_term
    dm_i .= dM
end

m = cu(ones(3, 10))  # Initial magnetization vector
dm = cu(ones(3, 10))
∂dm = cu(ones(3, 10))
∂m = cu(zeros(3, 10))
backend = get_backend(m)
bloch_kernel = bloch!(backend, 10, size(m, 2))
autodiff(Reverse, bloch_kernel, Duplicated(dm, ∂dm), Duplicated(m, ∂m), Const(params))[1]

The very MWE:

using KernelAbstractions
using CUDA
using DifferentiationInterface
using Enzyme

@kernel function emptyfunc()
    # dm = 0.0f0
end

empty_kernel = emptyfunc(CUDABackend(), 10, 10)
autodiff(Reverse, empty_kernel)

The error I get from both of these is:

ERROR: MethodError: no method matching iterate(::Nothing)
The function `iterate` exists, but no method is defined for this combination of argument types.

Closest candidates are:
  iterate(::Base.AsyncCollector, ::Base.AsyncCollectorState)
   @ Base asyncmap.jl:315
  iterate(::Base.AsyncCollector)
   @ Base asyncmap.jl:299
  iterate(::Pkg.Resolve.NodePerm, Any...)
   @ Pkg ~/julia/share/julia/stdlib/v1.11/Pkg/src/Resolve/maxsum.jl:240
  ...

Stacktrace:
  [1] _foldl_impl(op::Base.BottomRF{typeof(Base.mul_prod)}, init::Base._InitialValue, itr::Nothing)
    @ Base ./reduce.jl:56
  [2] foldl_impl(op::Base.BottomRF{typeof(Base.mul_prod)}, nt::Base._InitialValue, itr::Nothing)
    @ Base ./reduce.jl:48
  [3] mapfoldl_impl(f::typeof(identity), op::typeof(Base.mul_prod), nt::Base._InitialValue, itr::Nothing)
    @ Base ./reduce.jl:44
  [4] mapfoldl(f::Function, op::Function, itr::Nothing; init::Base._InitialValue)
    @ Base ./reduce.jl:175
  [5] mapreduce(f::Function, op::Function, itr::Nothing; kw::@Kwargs{})
    @ Base ./reduce.jl:307
  [6] prod(a::Nothing; kw::@Kwargs{})
    @ Base ./reduce.jl:617
  [7] _create_tape_kernel(::KernelAbstractions.Kernel{…}, ::ReverseModeSplit{…}, ::Type, ::Type, ::Nothing, ::KernelAbstractions.NDIteration.NDRange{…})
    @ EnzymeExt ~/.julia/packages/KernelAbstractions/HqVoT/ext/EnzymeCore08Ext.jl:107
  [8] macro expansion
    @ ~/.julia/packages/Enzyme/McaEh/src/compiler.jl:5610 [inlined]
  [9] enzyme_call
    @ ~/.julia/packages/Enzyme/McaEh/src/compiler.jl:5144 [inlined]
 [10] CombinedAdjointThunk
    @ ~/.julia/packages/Enzyme/McaEh/src/compiler.jl:5019 [inlined]
 [11] autodiff
    @ ~/.julia/packages/Enzyme/McaEh/src/Enzyme.jl:517 [inlined]
 [12] autodiff
    @ ~/.julia/packages/Enzyme/McaEh/src/Enzyme.jl:558 [inlined]
 [13] autodiff(::ReverseMode{…}, ::KernelAbstractions.Kernel{…})
    @ Enzyme ~/.julia/packages/Enzyme/McaEh/src/Enzyme.jl:530
 [14] top-level scope
    @ ~/KomaEnzyme/cuda_test.jl:75
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 with ext/EnzymeCore08Ext.jl:107 and reproduce the autodiff(Reverse, empty_kernel) MWE using CUDABackend(); compare it with the larger kernel example and the reported stack trace. Done means inability to differentiate is reported with a useful, intentional error rather than iterate(::Nothing), with coverage for both examples.

Written by the indexing model from the issue text.

Assessment

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.