Handle FIRK expanded AD matvec fallbacks

lib/BoundaryValueDiffEqFIRK/src/BoundaryValueDiffEqFIRK.jl

@@ -1,6 +1,6 @@
 module BoundaryValueDiffEqFIRK
 
-using ADTypes: ADTypes, AutoSparse, AutoForwardDiff
+using ADTypes: ADTypes, AutoSparse, AutoForwardDiff, AutoEnzyme, AutoMooncake
 using ArrayInterface: fast_scalar_indexing
 using BandedMatrices: BandedMatrix, Ones
 using BoundaryValueDiffEqCore: BoundaryValueDiffEqCore,

lib/BoundaryValueDiffEqFIRK/src/adaptivity.jl

@@ -201,17 +201,39 @@ function get_S_coeffs(h, yᵢ, yᵢ₊₁, dyᵢ, dyᵢ₊₁, ymid, dymid)
 end
 
 # S forward Interpolation
+function __firk_matvec!(y, A, b)
+    length(y) == size(A, 1) ||
+        throw(DimensionMismatch("y has length $(length(y)), but A has $(size(A, 1)) rows"))
+    length(b) == size(A, 2) ||
+        throw(DimensionMismatch("b has length $(length(b)), but A has $(size(A, 2)) columns"))
+
+    T = typeof(zero(eltype(A)) * zero(eltype(b)))
+    for (y_index, row) in zip(eachindex(y), axes(A, 1))
+        acc = zero(T)
+        for (b_index, col) in zip(eachindex(b), axes(A, 2))
+            @inbounds acc += A[row, col] * b[b_index]
+        end
+        @inbounds y[y_index] = acc
+    end
+    return y
+end
+
+function __firk_matvec(A, b)
+    y = similar(b, typeof(zero(eltype(A)) * zero(eltype(b))), size(A, 1))
+    return __firk_matvec!(y, A, b)
+end
+
 function S_interpolate!(y::AbstractArray, t, coeffs)
     ts = [t^(i - 1) for i in axes(coeffs, 2)]
-    return y .= coeffs * ts
+    return __firk_matvec!(y, coeffs, ts)
 end
 
 function dS_interpolate!(dy::AbstractArray, t, S_coeffs)
-    ts = zeros(size(S_coeffs, 2))
+    ts = zeros(promote_type(eltype(S_coeffs), typeof(t)), size(S_coeffs, 2))
     for i in 2:size(S_coeffs, 2)
         ts[i] = (i - 1) * t^(i - 2)
     end
-    return dy .= S_coeffs * ts
+    return __firk_matvec!(dy, S_coeffs, ts)
 end
 
 """
@@ -577,7 +599,7 @@ end
 end
 
 function get_q_coeffs(A, ki, h)
-    coeffs = A * ki
+    coeffs = __firk_matvec(A, ki)
     for i in axes(coeffs, 1)
         coeffs[i] = coeffs[i] / (h^(i - 1))
     end

lib/BoundaryValueDiffEqFIRK/src/collocation.jl

@@ -12,6 +12,11 @@ function Φ!(residual, cache::FIRKCacheNested, y, u, trait, constraint)
     )
 end
 
+@inline function __nested_stage_parameter(cache::FIRKCacheNested, p)
+    length(cache.bcresid_prototype) < cache.M && return nodual_value(p)
+    return p
+end
+
 @views function Φ!(
         residual, fᵢ_cache, k_discrete, f!, TU::FIRKTableau{false}, y, u, p,
         mesh, mesh_dt, stage::Int, f_prototype, singular_term, ::DiffCacheNeeded, ::Val{true}
@@ -195,7 +200,7 @@ end
 
         K = get_tmp(k_discrete[i], u)
 
-        _nestprob = remake(nest_prob, p = nestprob_p)
+        _nestprob = remake(nest_prob, p = __nested_stage_parameter(cache, nestprob_p))
         nestsol = __solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
         @. K = nestsol.u
         @. residᵢ = yᵢ₊₁ - yᵢ
@@ -227,7 +232,7 @@ end
 
         K = get_tmp(k_discrete[i], u)
 
-        _nestprob = remake(nest_prob, p = nestprob_p)
+        _nestprob = remake(nest_prob, p = __nested_stage_parameter(cache, nestprob_p))
         nestsol = __solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
         @. K = nestsol.u
         @. residᵢ = yᵢ₊₁ - yᵢ
@@ -385,7 +390,7 @@ end
         nestprob_p[2] = T(mesh_dt[i])
         nestprob_p[3:end] = yᵢ
 
-        _nestprob = remake(nest_prob, p = nestprob_p)
+        _nestprob = remake(nest_prob, p = __nested_stage_parameter(cache, nestprob_p))
         nestsol = __solve(_nestprob, nest_nlsolve_alg; alg.nested_nlsolve_kwargs...)
 
         @. residᵢ = yᵢ₊₁ - yᵢ

lib/BoundaryValueDiffEqFIRK/src/firk.jl

@@ -73,6 +73,33 @@ end
 
 Base.eltype(::FIRKCacheExpand{iip, T}) where {iip, T} = T
 
+# Julia 1.12+ exposes Enzyme/Mooncake failures in expanded FIRK collocation AD.
+const FIRK_EXPANDED_AD_FALLBACK = VERSION >= v"1.12.0-DEV.0"
+
+@inline __firk_expanded_diffmode(ad) = ad
+@inline function __firk_expanded_diffmode(ad::AutoEnzyme)
+    return FIRK_EXPANDED_AD_FALLBACK ? AutoForwardDiff() : ad
+end
+@inline function __firk_expanded_diffmode(ad::AutoMooncake)
+    return FIRK_EXPANDED_AD_FALLBACK ? AutoForwardDiff() : ad
+end
+@inline function __firk_expanded_diffmode(ad::AutoSparse)
+    dense_ad = __firk_expanded_diffmode(ADTypes.dense_ad(ad))
+    dense_ad === ADTypes.dense_ad(ad) && return ad
+    return AutoSparse(
+        dense_ad;
+        sparsity_detector = ADTypes.sparsity_detector(ad),
+        coloring_algorithm = ADTypes.coloring_algorithm(ad)
+    )
+end
+
+function __firk_expanded_jacobian_algorithm(jac_alg::BVPJacobianAlgorithm)
+    nonbc_diffmode = __firk_expanded_diffmode(jac_alg.nonbc_diffmode)
+    diffmode = __firk_expanded_diffmode(jac_alg.diffmode)
+    (nonbc_diffmode === jac_alg.nonbc_diffmode && diffmode === jac_alg.diffmode) && return jac_alg
+    return BVPJacobianAlgorithm(jac_alg.bc_diffmode, nonbc_diffmode, diffmode)
+end
+
 function extend_y(y, N::Int, stage::Int)
     y_extended = similar(y.u, (N - 1) * (stage + 1) + 1)
     for (i, ctr) in enumerate(2:(stage + 1):((N - 1) * (stage + 1) + 1))
@@ -294,6 +321,7 @@ function init_expanded(
     )
     verbose_spec = _process_verbose_param(verbose)
     @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg)
+    @set! alg.jac_alg = __firk_expanded_jacobian_algorithm(alg.jac_alg)
     iip = isinplace(prob)
     @assert (iip || isnothing(alg.optimize)) "Out-of-place constraints don't allow optimization solvers "
     if adaptive && isa(alg, FIRKNoAdaptivity)
@@ -1393,15 +1421,23 @@ function __construct_problem(
     )
 end
 
+@inline function __firk_eval_sol!(eval_sol::EvalSol, y_, cache)
+    u = __restructure_sol(y_, cache.in_size)
+    if eltype(first(u)) <: eltype(first(eval_sol.u))
+        eval_sol.u[1:end] .= u
+        return eval_sol
+    end
+    return EvalSol(u, eval_sol.t, cache)
+end
+
 @views function __firk_loss!(
         resid, u, p, y, pt::StandardBVProblem, bc!::BC, residual, mesh,
         cache, eval_sol, trait::DiffCacheNeeded, constraint
     ) where {BC}
     y_ = recursive_unflatten!(y, u)
     resids = [get_tmp(r, u) for r in residual]
     Φ!(resids[2:end], cache, y_, u, trait, constraint)
-    eval_sol.u[1:end] .= y_
-    eval_bc_residual!(resids[1], pt, bc!, eval_sol, p, mesh)
+    eval_bc_residual!(resids[1], pt, bc!, __firk_eval_sol!(eval_sol, y_, cache), p, mesh)
     recursive_flatten!(resid, resids)
     return nothing
 end
@@ -1413,8 +1449,7 @@ end
     y_ = recursive_unflatten!(y, u)
     resids = [r for r in residual]
     Φ!(resids[2:end], cache, y_, u, trait, constraint)
-    eval_sol.u[1:end] .= y_
-    eval_bc_residual!(resids[1], pt, bc!, eval_sol, p, mesh)
+    eval_bc_residual!(resids[1], pt, bc!, __firk_eval_sol!(eval_sol, y_, cache), p, mesh)
     recursive_flatten!(resid, resids)
     return nothing
 end
@@ -1473,8 +1508,7 @@ end
         u, p, y, pt::StandardBVProblem, bc::BC, mesh, cache, eval_sol, trait
     ) where {BC}
     y_ = recursive_unflatten!(y, u)
-    eval_sol.u[1:end] .= y_
-    resid_bc = eval_bc_residual(pt, bc, eval_sol, p, mesh)
+    resid_bc = eval_bc_residual(pt, bc, __firk_eval_sol!(eval_sol, y_, cache), p, mesh)
     resid_co = Φ(cache, y_, u, trait)
     return vcat(resid_bc, mapreduce(vec, vcat, resid_co))
 end
@@ -1522,7 +1556,10 @@ end
         resid, u, p, y, mesh, residual, cache, trait::NoDiffCacheNeeded, constraint
     )
     y_ = recursive_unflatten!(y, u)
-    resids = [r for r in residual[2:end]]
+    resids = Vector{eltype(residual)}(undef, length(residual) - 1)
+    for i in eachindex(resids)
+        resids[i] = residual[i + 1]
+    end
     Φ!(resids, cache, y_, u, trait, constraint)
     recursive_flatten!(resid, resids)
     return nothing

lib/BoundaryValueDiffEqFIRK/src/interpolation.jl

@@ -386,7 +386,7 @@ function (s::EvalSol{C})(tval::Number) where {C <: FIRKCacheExpand}
 end
 
 function get_q_coeffs_interp(A, ki, h)
-    coeffs = A * ki
+    coeffs = __firk_matvec(A, ki)
     for i in axes(coeffs, 1)
         coeffs[i] = coeffs[i] / (h^(i - 1))
     end