Handle FIRK prerelease AD matvec fallbacks

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

Author: ChrisRackauckas

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -100,6 +100,26 @@ function __maybe_matmul!(z::AbstractArray, A, b, α = eltype(z)(1), β = eltype(
     return mul!(z, A, b, α, β)
 end
 
+function __maybe_matmul!(
+        z::AbstractVector{<:Number}, A::AbstractMatrix{<:Number},
+        b::AbstractVector{<:Number}, α = eltype(z)(1), β = eltype(z)(0)
+    )
+    length(z) == size(A, 1) ||
+        throw(DimensionMismatch("z has length $(length(z)), 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 = promote_type(eltype(z), eltype(A), eltype(b), typeof(α), typeof(β))
+    for (z_index, row) in zip(eachindex(z), 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 z[z_index] = α * acc + β * z[z_index]
+    end
+    return z
+end
+
 # NOTE: We can implement it as mul! as above but then we pay the cost of moving
 #       `w` to the GPU too many times. Instead if we iterate of w and w′ we save
 #       that cost. Our main cost is anyways going to be due to a large `u0` and

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 = promote_type(eltype(y), eltype(A), 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, promote_type(eltype(A), 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/firk.jl

@@ -73,6 +73,34 @@ end
 
 Base.eltype(::FIRKCacheExpand{iip, T}) where {iip, T} = T
 
+# Julia 1.13 prereleases expose Enzyme/Mooncake issues in expanded FIRK collocation AD.
+const FIRK_AD_PRERELEASE_FALLBACK =
+    VERSION >= v"1.13.0-DEV.0" && VERSION < v"1.13.0"
+
+@inline __firk_prerelease_diffmode(ad) = ad
+@inline function __firk_prerelease_diffmode(ad::AutoEnzyme)
+    return FIRK_AD_PRERELEASE_FALLBACK ? AutoForwardDiff() : ad
+end
+@inline function __firk_prerelease_diffmode(ad::AutoMooncake)
+    return FIRK_AD_PRERELEASE_FALLBACK ? AutoForwardDiff() : ad
+end
+@inline function __firk_prerelease_diffmode(ad::AutoSparse)
+    dense_ad = __firk_prerelease_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_prerelease_jacobian_algorithm(jac_alg::BVPJacobianAlgorithm)
+    nonbc_diffmode = __firk_prerelease_diffmode(jac_alg.nonbc_diffmode)
+    diffmode = __firk_prerelease_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 +322,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_prerelease_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)

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