Fix MIRK AD fallback and view flattening

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

Author: ChrisRackauckas

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -1,24 +1,24 @@
 _unwrap_val(::Val{B}) where {B} = B
 _unwrap_val(B) = B
 
-recursive_length(x::Vector{<:AbstractArray}) = sum(length, x)
+recursive_length(x::AbstractVector{<:AbstractArray}) = sum(length, x)
 recursive_length(x::Vector{<:DiffCache}) = sum(xᵢ -> length(xᵢ.u), x)
 
-function recursive_flatten(x::Vector{<:AbstractArray})
+function recursive_flatten(x::AbstractVector{<:AbstractArray})
     y = zero(first(x), recursive_length(x))
     recursive_flatten!(y, x)
     return y
 end
 
-@views function recursive_flatten!(y::AbstractVector, x::Vector{<:AbstractArray})
+@views function recursive_flatten!(y::AbstractVector, x::AbstractVector{<:AbstractArray})
     i = 0
     for xᵢ in x
         copyto!(y[(i + 1):(i + length(xᵢ))], xᵢ)
         i += length(xᵢ)
     end
     return y
 end
-@views function recursive_flatten_twopoint!(y::AbstractVector, x::Vector{<:AbstractArray}, sizes)
+@views function recursive_flatten_twopoint!(y::AbstractVector, x::AbstractVector{<:AbstractArray}, sizes)
     x_, xiter = first(x), x[2:end]
     copyto!(y[1:prod(sizes[1])], x_[1:prod(sizes[1])])
     i = prod(sizes[1])
@@ -30,7 +30,7 @@ end
     return y
 end
 
-@views function recursive_unflatten!(y::Vector{<:AbstractArray}, x::AbstractVector)
+@views function recursive_unflatten!(y::AbstractVector{<:AbstractArray}, x::AbstractVector)
     i = 0
     for yᵢ in y
         copyto!(yᵢ, x[(i + 1):(i + length(yᵢ))])

lib/BoundaryValueDiffEqMIRK/src/BoundaryValueDiffEqMIRK.jl

@@ -1,6 +1,6 @@
 module BoundaryValueDiffEqMIRK
 
-using ADTypes: ADTypes, AutoForwardDiff, AutoSparse
+using ADTypes: ADTypes, AutoForwardDiff, AutoSparse, AutoEnzyme, AutoMooncake
 using ArrayInterface: fast_scalar_indexing
 using BandedMatrices: BandedMatrix, Ones
 using BoundaryValueDiffEqCore: BoundaryValueDiffEqCore,
@@ -45,7 +45,7 @@ using SciMLBase: SciMLBase, AbstractDiffEqInterpolation, BVPFunction, BVProblem,
     __solve, isinplace, remake, solve
 using Setfield: @set!
 using Reexport: @reexport
-using PreallocationTools: PreallocationTools, get_tmp
+using PreallocationTools: PreallocationTools, DiffCache, get_tmp
 using PrecompileTools: @compile_workload, @setup_workload
 using Preferences: Preferences
 using SparseArrays: sparse

lib/BoundaryValueDiffEqMIRK/src/interpolation.jl

@@ -224,8 +224,7 @@ function (s::EvalSol{C})(tval::Number) where {C <: MIRKCache}
     dt = cache.mesh_dt[ii]
     τ = (tval - t[ii]) / dt
     w, _ = interp_weights(τ, alg)
-    K = __needs_diffcache(alg.jac_alg) ? @view(k_discrete[ii].du[:, 1:stage]) :
-        @view(k_discrete[ii][:, 1:stage])
+    K = __mirk_discrete_stage_view(cache, ii, stage)
     KI = @view(k_interp.u[ii][1:length_z, 1:(cache.ITU.s_star - stage)])
     __maybe_matmul!(@view(z[1:length_z]), K, @view(w[1:stage]))
     __maybe_matmul!(@view(z[1:length_z]), KI, @view(w[(stage + 1):cache.ITU.s_star]), true, true)
@@ -256,8 +255,7 @@ function (s::EvalSol{C})(tvals::AbstractArray{<:Number}) where {C <: MIRKCache}
         dt = mesh_dt[ii]
         τ = (tval - t[ii]) / dt
         w, _ = interp_weights(τ, alg)
-        K = __needs_diffcache(alg.jac_alg) ? @view(k_discrete[ii].du[:, 1:stage]) :
-            @view(k_discrete[ii][:, 1:stage])
+        K = __mirk_discrete_stage_view(cache, ii, stage)
         KI = @view(k_interp.u[ii][1:length_z, 1:(cache.ITU.s_star - stage)])
         __maybe_matmul!(@view(zvals[i][1:length_z]), K, @view(w[1:stage]))
         __maybe_matmul!(
@@ -283,14 +281,26 @@ function (s::EvalSol{C})(tval::Number, ::Type{Val{1}}) where {C <: MIRKCache}
     dt = mesh_dt[ii]
     τ = (tval - t[ii]) / dt
     _, w′ = interp_weights(τ, alg)
-    __maybe_matmul!(z′, @view(k_discrete[ii].du[:, 1:stage]), @view(w′[1:stage]))
+    __maybe_matmul!(z′, __mirk_discrete_stage_view(cache, ii, stage), @view(w′[1:stage]))
     __maybe_matmul!(
         z′, @view(k_interp.u[ii][:, 1:(cache.ITU.s_star - stage)]), @view(w′[(stage + 1):cache.ITU.s_star]),
         true, true
     )
     return z′
 end
 
+@inline function __mirk_discrete_stage_view(
+        cache::MIRKCache{iip, T, use_both, DiffCacheNeeded}, ii, stage
+    ) where {iip, T, use_both}
+    return @view(cache.k_discrete[ii].du[:, 1:stage])
+end
+
+@inline function __mirk_discrete_stage_view(
+        cache::MIRKCache{iip, T, use_both, NoDiffCacheNeeded}, ii, stage
+    ) where {iip, T, use_both}
+    return @view(cache.k_discrete[ii][:, 1:stage])
+end
+
 """
     interp_setup!(cache::MIRKCache)
 

lib/BoundaryValueDiffEqMIRK/src/mirk.jl

@@ -51,7 +51,7 @@ function SciMLBase.__init(
     verbose_spec = _process_verbose_param(verbose)
     @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg)
     iip = isinplace(prob)
-    diffcache = __cache_trait(alg.jac_alg)
+    diffcache = __mirk_iip_cache_trait(Val(iip), alg.jac_alg)
     @assert (iip || isnothing(alg.optimize)) "Out-of-place constraints don't allow optimization solvers "
 
     tune_parameters = haskey(prob.kwargs, :tune_parameters)
@@ -74,7 +74,9 @@ function SciMLBase.__init(
     mesh_dt = diff(mesh)
 
     chunksize = pickchunksize(N * (Nig - 1))
-    __alloc = @closure x -> __maybe_allocate_diffcache(vec(zero(x)), chunksize, alg.jac_alg)
+    __alloc = @closure x -> __mirk_iip_maybe_allocate_diffcache(
+        vec(zero(x)), chunksize, Val(iip), alg.jac_alg
+    )
 
     fᵢ_cache = __alloc(zero(u0))
     fᵢ₂_cache = vec(zero(u0))
@@ -91,12 +93,16 @@ function SciMLBase.__init(
 
     k_discrete = if !constraint
         [
-            __maybe_allocate_diffcache(safe_similar(u0, N, stage), chunksize, alg.jac_alg)
+            __mirk_iip_maybe_allocate_diffcache(
+                    safe_similar(u0, N, stage), chunksize, Val(iip), alg.jac_alg
+                )
                 for _ in 1:Nig
         ]
     else
         [
-            __maybe_allocate_diffcache(safe_similar(u0, L_f_prototype, stage), chunksize, alg.jac_alg)
+            __mirk_iip_maybe_allocate_diffcache(
+                    safe_similar(u0, L_f_prototype, stage), chunksize, Val(iip), alg.jac_alg
+                )
                 for _ in 1:Nig
         ]
     end
@@ -391,6 +397,33 @@ function __construct_problem(cache::MIRKCache{iip}, y::AbstractVector, y₀::Abs
     return __construct_problem(cache, y, y₀, Val(constraint))
 end
 
+__mirk_iip_ad_diffmode(::Val{true}, ::AutoEnzyme) = AutoForwardDiff()
+__mirk_iip_ad_diffmode(::Val{true}, ::AutoMooncake) = AutoForwardDiff()
+function __mirk_iip_ad_diffmode(iip, diffmode::AutoSparse)
+    return AutoSparse(
+        __mirk_iip_ad_diffmode(iip, get_dense_ad(diffmode));
+        sparsity_detector = __default_sparsity_detector(diffmode),
+        coloring_algorithm = __default_coloring_algorithm(diffmode)
+    )
+end
+__mirk_iip_ad_diffmode(_, diffmode) = diffmode
+
+__mirk_iip_needs_diffcache(iip, diffmode::AutoSparse) = __mirk_iip_needs_diffcache(
+    iip, get_dense_ad(diffmode)
+)
+__mirk_iip_needs_diffcache(iip, diffmode) = __needs_diffcache(__mirk_iip_ad_diffmode(iip, diffmode))
+function __mirk_iip_needs_diffcache(iip, jac_alg::BVPJacobianAlgorithm)
+    return __mirk_iip_needs_diffcache(iip, jac_alg.diffmode) ||
+        __mirk_iip_needs_diffcache(iip, jac_alg.bc_diffmode) ||
+        __mirk_iip_needs_diffcache(iip, jac_alg.nonbc_diffmode)
+end
+__mirk_iip_cache_trait(iip, jac_alg) = __mirk_iip_needs_diffcache(iip, jac_alg) ?
+    DiffCacheNeeded() : NoDiffCacheNeeded()
+function __mirk_iip_maybe_allocate_diffcache(x, chunksize, iip, jac_alg)
+    return __mirk_iip_needs_diffcache(iip, jac_alg) ?
+        DiffCache(x, chunksize; warn_on_resize = false) : x
+end
+
 function __construct_problem(
         cache::MIRKCache{iip}, y::AbstractVector,
         y₀::AbstractVectorOfArray, constraint
@@ -400,7 +433,7 @@ function __construct_problem(
 
     eval_sol = EvalSol(__restructure_sol(y₀.u, cache.in_size), cache.mesh, cache)
 
-    trait = __cache_trait(jac_alg)
+    trait = __mirk_iip_cache_trait(Val(iip), jac_alg)
 
     loss_bc = if iip
         @closure (
@@ -614,7 +647,7 @@ function __construct_problem(
     ) where {iip, T, UB, DC, tune_parameters, BC, C, LF}
     (; jac_alg) = cache.alg
     (; f_prototype, bcresid_prototype, prob) = cache
-    (; bc_diffmode) = jac_alg
+    bc_diffmode = __mirk_iip_ad_diffmode(Val(iip), jac_alg.bc_diffmode)
     N = length(cache.mesh)
 
     resid_bc = bcresid_prototype
@@ -623,7 +656,8 @@ function __construct_problem(
     resid_collocation = safe_similar(y, L_f_prototype * (N - 1))
     loss_bc_ad = if iip
         @closure (u, p) -> __mirk_loss_bc_iip_ad(
-            u, p, StandardBVProblem(), cache.bc, cache.y, cache.mesh, cache, __cache_trait(jac_alg)
+            u, p, StandardBVProblem(), cache.bc, cache.y, cache.mesh, cache,
+            __mirk_iip_cache_trait(Val(iip), jac_alg)
         )
     else
         loss_bc
@@ -634,7 +668,7 @@ function __construct_problem(
     )
 
     nonbc_diffmode = AutoSparse(
-        get_dense_ad(jac_alg.nonbc_diffmode),
+        __mirk_iip_ad_diffmode(Val(iip), get_dense_ad(jac_alg.nonbc_diffmode)),
         sparsity_detector = __default_sparsity_detector(jac_alg.nonbc_diffmode),
         coloring_algorithm = __default_coloring_algorithm(jac_alg.nonbc_diffmode)
     )
@@ -700,7 +734,7 @@ function __construct_problem(
     ) where {iip, T, UB, DC, tune_parameters, BC, C, LF}
     (; jac_alg) = cache.alg
     (; f_prototype, bcresid_prototype, prob) = cache
-    (; bc_diffmode) = jac_alg
+    bc_diffmode = __mirk_iip_ad_diffmode(Val(iip), jac_alg.bc_diffmode)
     N = length(cache.mesh)
 
     resid_bc = bcresid_prototype
@@ -709,7 +743,8 @@ function __construct_problem(
     resid_prototype = vcat(resid_bc, resid_collocation)
     loss_bc_ad = if iip
         @closure (u, p) -> __mirk_loss_bc_iip_ad(
-            u, p, StandardBVProblem(), cache.bc, cache.y, cache.mesh, cache, __cache_trait(jac_alg)
+            u, p, StandardBVProblem(), cache.bc, cache.y, cache.mesh, cache,
+            __mirk_iip_cache_trait(Val(iip), jac_alg)
         )
     else
         loss_bc
@@ -736,13 +771,13 @@ function __construct_problem(
             )
         end
         AutoSparse(
-            get_dense_ad(jac_alg.nonbc_diffmode);
+            __mirk_iip_ad_diffmode(Val(iip), get_dense_ad(jac_alg.nonbc_diffmode));
             sparsity_detector = ADTypes.KnownJacobianSparsityDetector(sparse_jacobian_prototype),
             coloring_algorithm = __default_coloring_algorithm(jac_alg.nonbc_diffmode)
         )
     else
         J_full_band = nothing
-        jac_alg.nonbc_diffmode
+        __mirk_iip_ad_diffmode(Val(iip), jac_alg.nonbc_diffmode)
     end
 
     cache_collocation = if iip
@@ -903,12 +938,12 @@ function __construct_problem(
 
     diffmode = if jac_alg.diffmode isa AutoSparse
         AutoSparse(
-            get_dense_ad(jac_alg.diffmode);
+            __mirk_iip_ad_diffmode(Val(iip), get_dense_ad(jac_alg.diffmode));
             sparsity_detector = __default_sparsity_detector(jac_alg.diffmode),
             coloring_algorithm = __default_coloring_algorithm(jac_alg.diffmode)
         )
     else
-        jac_alg.diffmode
+        __mirk_iip_ad_diffmode(Val(iip), jac_alg.diffmode)
     end
 
     diffcache = if iip
@@ -970,12 +1005,12 @@ function __construct_problem(
             @view(bcresid_prototype[(prod(cache.resid_size[1]) + 1):end]), cache.M, N
         )
         AutoSparse(
-            get_dense_ad(jac_alg.diffmode);
+            __mirk_iip_ad_diffmode(Val(iip), get_dense_ad(jac_alg.diffmode));
             sparsity_detector = ADTypes.KnownJacobianSparsityDetector(sparse_jacobian_prototype),
             coloring_algorithm = __default_coloring_algorithm(jac_alg.diffmode)
         )
     else
-        jac_alg.diffmode
+        __mirk_iip_ad_diffmode(Val(iip), jac_alg.diffmode)
     end
 
     diffcache = if iip