Implicit methods don't work with `DynamicalODEProblem` on GPUs

[efaulhaber]

MWE:

using OrdinaryDiffEqSDIRK, ADTypes, CUDA

function f1(dv, v, u, p, t)
    dv .= 0
    return dv
end

function f2(du, v, u, p, t)
    du .= 0
    return du
end

u0 = CUDA.zeros(2)
v0 = CUDA.zeros(2)

tspan = (0.0, 1.0)
prob = DynamicalODEProblem(f1, f2, v0, u0, tspan, nothing)

sol = solve(prob, KenCarp3(autodiff=AutoFiniteDiff()), save_everystep = false)

ERROR: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/uiVyU/src/host/indexing.jl:50 [inlined]
  [6] _broadcast_getindex
    @ ./broadcast.jl:644 [inlined]
  [7] _getindex
    @ ./broadcast.jl:675 [inlined]
  [8] _broadcast_getindex
    @ ./broadcast.jl:650 [inlined]
  [9] getindex
    @ ./broadcast.jl:610 [inlined]
 [10] cpu_broadcast_kernel_linear
    @ ~/.julia/packages/KernelAbstractions/sWSE0/src/macros.jl:304 [inlined]
 [11] (::GPUArrays.var"#cpu_broadcast_kernel_linear#37")(__ctx__::KernelAbstractions.CompilerMetadata{…}, dest::SubArray{…}, bc::Base.Broadcast.Broadcasted{…})
    @ GPUArrays ./none:0
 [12] __thread_run(tid::Int64, len::Int64, rem::Int64, obj::KernelAbstractions.Kernel{…}, ndrange::Tuple{…}, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.NoDynamicCheck)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/sWSE0/src/cpu.jl:144
 [13] __run(obj::KernelAbstractions.Kernel{…}, ndrange::Tuple{…}, iterspace::KernelAbstractions.NDIteration.NDRange{…}, args::Tuple{…}, dynamic::KernelAbstractions.NDIteration.NoDynamicCheck, static_threads::Bool)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/sWSE0/src/cpu.jl:111
 [14] (::KernelAbstractions.Kernel{…})(::SubArray{…}, ::Vararg{…}; ndrange::Tuple{…}, workgroupsize::Nothing)
    @ KernelAbstractions ~/.julia/packages/KernelAbstractions/sWSE0/src/cpu.jl:45
 [15] _copyto!
    @ ~/.julia/packages/GPUArrays/uiVyU/src/host/broadcast.jl:71 [inlined]
 [16] materialize!
    @ ~/.julia/packages/GPUArrays/uiVyU/src/host/broadcast.jl:38 [inlined]
 [17] materialize!
    @ ./broadcast.jl:880 [inlined]
 [18] copyto!(dest::Vector{Float32}, A::RecursiveArrayTools.ArrayPartition{Float32, Tuple{CuArray{…}, CuArray{…}}})
    @ RecursiveArrayTools ~/.julia/packages/RecursiveArrayTools/iAeBW/src/array_partition.jl:187
 [19] copyto_axcheck!
    @ ./abstractarray.jl:1167 [inlined]
 [20] Array
    @ ./array.jl:626 [inlined]
 [21] Array
    @ ./boot.jl:603 [inlined]
 [22] zeromatrix(A::RecursiveArrayTools.ArrayPartition{Float32, Tuple{CuArray{…}, CuArray{…}}})
    @ RecursiveArrayTools ~/.julia/packages/RecursiveArrayTools/iAeBW/src/array_partition.jl:433
 [23] build_J_W(alg::KenCarp3{…}, u::RecursiveArrayTools.ArrayPartition{…}, uprev::RecursiveArrayTools.ArrayPartition{…}, p::Nothing, t::Float64, dt::Float64, f::DynamicalODEFunction{…}, jac_config::Tuple{…}, ::Type{…}, ::Val{…})
    @ OrdinaryDiffEqDifferentiation ~/.julia/packages/OrdinaryDiffEqDifferentiation/OY5R1/src/derivative_utils.jl:836
 [24] build_nlsolver(alg::KenCarp3{…}, nlalg::OrdinaryDiffEqNonlinearSolve.NLNewton{…}, u::RecursiveArrayTools.ArrayPartition{…}, uprev::RecursiveArrayTools.ArrayPartition{…}, p::Nothing, t::Float64, dt::Float64, f::DynamicalODEFunction{…}, rate_prototype::RecursiveArrayTools.ArrayPartition{…}, ::Type{…}, ::Type{…}, ::Type{…}, γ::Float64, c::Float64, α::Int64, ::Val{…})
    @ OrdinaryDiffEqNonlinearSolve ~/.julia/packages/OrdinaryDiffEqNonlinearSolve/6Gmae/src/utils.jl:201
 [25] build_nlsolver
    @ ~/.julia/packages/OrdinaryDiffEqNonlinearSolve/6Gmae/src/utils.jl:152 [inlined]
 [26] build_nlsolver
    @ ~/.julia/packages/OrdinaryDiffEqNonlinearSolve/6Gmae/src/utils.jl:142 [inlined]
 [27] alg_cache(alg::KenCarp3{…}, u::RecursiveArrayTools.ArrayPartition{…}, rate_prototype::RecursiveArrayTools.ArrayPartition{…}, ::Type{…}, ::Type{…}, ::Type{…}, uprev::RecursiveArrayTools.ArrayPartition{…}, uprev2::RecursiveArrayTools.ArrayPartition{…}, f::DynamicalODEFunction{…}, t::Float64, dt::Float64, reltol::Float32, p::Nothing, calck::Bool, ::Val{…})
    @ OrdinaryDiffEqSDIRK ~/.julia/packages/OrdinaryDiffEqSDIRK/YAA9M/src/kencarp_kvaerno_caches.jl:96
 [28] __init(prob::ODEProblem{…}, alg::KenCarp3{…}, timeseries_init::Tuple{}, ts_init::Tuple{}, ks_init::Tuple{}, recompile::Type{…}; saveat::Tuple{}, tstops::Tuple{}, d_discontinuities::Tuple{}, save_idxs::Nothing, save_everystep::Bool, save_on::Bool, save_start::Bool, save_end::Nothing, callback::Nothing, dense::Bool, calck::Bool, dt::Float64, dtmin::Float64, dtmax::Float64, force_dtmin::Bool, adaptive::Bool, gamma::Rational{…}, abstol::Nothing, reltol::Nothing, qmin::Rational{…}, qmax::Int64, qsteady_min::Int64, qsteady_max::Rational{…}, beta1::Nothing, beta2::Nothing, qoldinit::Rational{…}, controller::Nothing, fullnormalize::Bool, failfactor::Int64, maxiters::Int64, internalnorm::typeof(DiffEqBase.ODE_DEFAULT_NORM), internalopnorm::typeof(LinearAlgebra.opnorm), isoutofdomain::typeof(DiffEqBase.ODE_DEFAULT_ISOUTOFDOMAIN), unstable_check::typeof(DiffEqBase.ODE_DEFAULT_UNSTABLE_CHECK), verbose::Bool, timeseries_errors::Bool, dense_errors::Bool, advance_to_tstop::Bool, stop_at_next_tstop::Bool, initialize_save::Bool, progress::Bool, progress_steps::Int64, progress_name::String, progress_message::typeof(DiffEqBase.ODE_DEFAULT_PROG_MESSAGE), progress_id::Symbol, userdata::Nothing, allow_extrapolation::Bool, initialize_integrator::Bool, alias::ODEAliasSpecifier, initializealg::OrdinaryDiffEqCore.DefaultInit, kwargs::@Kwargs{})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/yWlnm/src/solve.jl:409
 [29] __init (repeats 5 times)
    @ ~/.julia/packages/OrdinaryDiffEqCore/yWlnm/src/solve.jl:11 [inlined]
 [30] __solve(::ODEProblem{…}, ::KenCarp3{…}; kwargs::@Kwargs{…})
    @ OrdinaryDiffEqCore ~/.julia/packages/OrdinaryDiffEqCore/yWlnm/src/solve.jl:6
 [31] __solve
    @ ~/.julia/packages/OrdinaryDiffEqCore/yWlnm/src/solve.jl:1 [inlined]
 [32] #solve_call#36
    @ ~/.julia/packages/DiffEqBase/ajV2I/src/solve.jl:667 [inlined]
 [33] solve_call
    @ ~/.julia/packages/DiffEqBase/ajV2I/src/solve.jl:624 [inlined]
 [34] #solve_up#45
    @ ~/.julia/packages/DiffEqBase/ajV2I/src/solve.jl:1199 [inlined]
 [35] solve_up
    @ ~/.julia/packages/DiffEqBase/ajV2I/src/solve.jl:1177 [inlined]
 [36] #solve#43
    @ ~/.julia/packages/DiffEqBase/ajV2I/src/solve.jl:1089 [inlined]
 [37] top-level scope
    @ ~/git/PointNeighbors.jl/benchmarks/test_gpu.jl:19
Some type information was truncated. Use `show(err)` to see complete types.

[ChrisRackauckas]

The methods which specialize on DynamicalODEProblem should already work, but it seems like the ODEProblem fallbacks can have a few issues. I solved some of them in SciML/RecursiveArrayTools.jl#444.

The remaining part is the Jacobian construction which is a bit harder.