Improve error on Enzyme inability to differentiate

[KookiesNKareem]

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.

[KookiesNKareem]

The fix was pretty simple; it couldn't iterate because I didn't set the ndrange parameter of the Kernel properly.

[vchuravy]

That should throw a better error....

[KookiesNKareem]

Well this fix caused an Enzyme internal error, which told me to open an issue so I think I have to open a different issue for that

[KookiesNKareem]

Sorry didn't mean to close

[cwittens]

That should throw a better error....

I just had a similar problem (when not using Enzyme), where I also forgot to include the ndrange and got the following error:

on CPU() backend

ERROR: MethodError: no method matching iterate(::Nothing)

or on ROCBackend() backend:

ERROR: MethodError: no method matching length(::Nothing)