Improve error on Enzyme inability to differentiate #629
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.