Eliminate final per-step allocation via y_cache and add QA allocation test

Add y_cache field to MIRKCache that pre-allocates the Vector{Vector{T}}
needed by recursive_unflatten! for the primal (Float64) path. This
eliminates the last 144 bytes/call from get_tmp. broadcast allocation.

For the Dual path (ForwardDiff Jacobian computation), falls back to the
existing broadcast since element types differ.

Loss function is now fully allocation-free at runtime (0 bytes/call
verified via @timed over 100k calls).

Add zero-allocation QA test in the qa test group.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: ChrisRackauckas

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -68,6 +68,28 @@ end
     return recursive_unflatten!(get_tmp.(y, (x,)), x)
 end
 
+# Non-allocating version with pre-allocated output cache.
+# When element types match (primal path), fills y_cache in-place.
+# When they don't (Dual path), falls back to broadcast allocation.
+@views function recursive_unflatten!(
+        y::Vector{<:DiffCache}, y_cache::Vector{<:AbstractVector{T}}, x::AbstractVector{T}
+    ) where {T}
+    i = 0
+    for (j, yᵢ) in enumerate(y)
+        tmp = PreallocationTools.get_tmp(yᵢ, x)
+        y_cache[j] = tmp
+        copyto!(tmp, x[(i + 1):(i + length(tmp))])
+        i += length(tmp)
+    end
+    return y_cache
+end
+
+@views function recursive_unflatten!(
+        y::Vector{<:DiffCache}, y_cache::Vector, x::AbstractVector
+    )
+    return recursive_unflatten!(get_tmp.(y, (x,)), x)
+end
+
 @views function recursive_unflatten!(y::AbstractVectorOfArray, x::AbstractVector)
     i = 0
     for yᵢ in y

lib/BoundaryValueDiffEqMIRK/src/mirk.jl

@@ -22,9 +22,10 @@
     y
     y₀
     residual
-    # The following 2 caches are never resized
+    # The following 3 caches are never resized
     fᵢ_cache
     fᵢ₂_cache
+    y_cache                        # Pre-allocated get_tmp pointers for primal path
     errors
     new_stages
     resid_size
@@ -66,6 +67,7 @@ function SciMLBase.__init(
 
     fᵢ_cache = __alloc(zero(X))
     fᵢ₂_cache = vec(zero(X))
+    y_cache = Vector{Vector{T}}(undef, Nig + 1)
 
     # Don't flatten this here, since we need to expand it later if needed
     y₀ = __initial_guess_on_mesh(X, mesh, prob.p)
@@ -231,7 +233,7 @@ function SciMLBase.__init(
     return MIRKCache{iip, T, use_both, typeof(diffcache), tune_parameters}(
         alg_order(alg), stage, N, size(X), f, bc, prob_, prob.problem_type, prob.p, alg,
         TU, ITU, f_prototype, bcresid_prototype, mesh, mesh_dt, k_discrete, k_interp, y,
-        y₀, residual, fᵢ_cache, fᵢ₂_cache, errors, new_stages, resid₁_size, prob.singular_term
+        y₀, residual, fᵢ_cache, fᵢ₂_cache, y_cache, errors, new_stages, resid₁_size, prob.singular_term
         , nlsolve_kwargs, optimize_kwargs, (; abstol, dt, adaptive, controller, tune_parameters, kwargs...), verbose_spec
     )
 end
@@ -251,6 +253,7 @@ function __expand_cache!(cache::MIRKCache{iip, T, use_both}) where {iip, T, use_
     __resize!(cache.residual, Nₙ, cache.M)
     __resize!(cache.errors, ifelse(use_both, 2 * (Nₙ - 1), (Nₙ - 1)), cache.M)
     __resize!(cache.new_stages, Nₙ - 1, cache.M)
+    resize!(cache.y_cache, Nₙ)
     return cache
 end
 
@@ -439,8 +442,8 @@ end
         resid, u, p, y, pt::StandardBVProblem, bc!::BC, residual, mesh,
         cache, EvalSol, trait::DiffCacheNeeded, constraint
     ) where {BC}
-    y_ = recursive_unflatten!(y, u)
-    Φ!(residual[2:end], cache, y_, u, trait, constraint)
+    y_ = recursive_unflatten!(y, cache.y_cache, u)
+    Φ!(residual[2:end], cache, y, u, trait, constraint)
     copyto!(EvalSol.u, __restructure_sol(y_, cache.in_size))
     copyto!(EvalSol.cache.k_discrete, cache.k_discrete)
     eval_bc_residual!(get_tmp(residual[1], u), pt, bc!, EvalSol, p, mesh)
@@ -478,8 +481,8 @@ end
         resid, u, p, y, pt::TwoPointBVProblem, bc!::Tuple{BC1, BC2}, residual,
         mesh, cache, _, trait::DiffCacheNeeded, constraint
     ) where {BC1, BC2}
-    y_ = recursive_unflatten!(y, u)
-    Φ!(residual[2:end], cache, y_, u, trait, constraint)
+    y_ = recursive_unflatten!(y, cache.y_cache, u)
+    Φ!(residual[2:end], cache, y, u, trait, constraint)
     resid0 = get_tmp(residual[1], u)
     resida = resid0[1:prod(cache.resid_size[1])]
     residb = resid0[(prod(cache.resid_size[1]) + 1):end]
@@ -551,9 +554,9 @@ end
 @views function __mirk_loss_collocation!(
         resid, u, p, y, mesh, residual, cache, trait::DiffCacheNeeded, constraint
     )
-    y_ = recursive_unflatten!(y, u)
+    recursive_unflatten!(y, cache.y_cache, u)
     collocation_residual = residual[2:end]
-    Φ!(collocation_residual, cache, y_, u, trait, constraint)
+    Φ!(collocation_residual, cache, y, u, trait, constraint)
     recursive_flatten!(resid, collocation_residual, u)
     return nothing
 end

test/qa/runtests.jl

@@ -1,4 +1,5 @@
-using BoundaryValueDiffEq, Aqua, JET, Test, SciMLBase
+using BoundaryValueDiffEq, BoundaryValueDiffEqMIRK, BoundaryValueDiffEqCore,
+    Aqua, JET, Test, SciMLBase
 
 @testset "Quality Assurance" begin
     @testset "Aqua" begin
@@ -15,4 +16,40 @@ using BoundaryValueDiffEq, Aqua, JET, Test, SciMLBase
         )
         @test length(JET.get_reports(rep)) == 0
     end
+
+    @testset "Zero per-step allocations in MIRK loss function" begin
+        function _f!(du, u, p, t)
+            du[1] = u[2]
+            du[2] = -u[1]
+            return nothing
+        end
+        function _bc!(resid, sol, p, t)
+            resid[1] = sol(0.0)[1] - 1.0
+            resid[2] = sol(1.0)[1] - cos(1.0)
+            return nothing
+        end
+        u0 = [1.0, 0.0]
+        tspan = (0.0, 1.0)
+        bvp = BVProblem(
+            BVPFunction{true}(_f!, _bc!; bcresid_prototype = zeros(2)), u0, tspan)
+
+        cache = SciMLBase.__init(bvp, MIRK4(); dt = 0.1, adaptive = false)
+        nlprob = BoundaryValueDiffEqMIRK.__construct_problem(
+            cache, vec(cache.y₀), copy(cache.y₀))
+
+        u_test = copy(nlprob.u0)
+        resid_test = zeros(length(nlprob.u0))
+        p_test = nlprob.p
+
+        # Verify loss function is allocation-free at runtime
+        function _bench_loss(f, resid, u, p, N)
+            for _ in 1:N
+                f(resid, u, p)
+            end
+        end
+        _bench_loss(nlprob.f, resid_test, u_test, p_test, 10)  # warmup
+        stats = @timed _bench_loss(nlprob.f, resid_test, u_test, p_test, 10_000)
+        bytes_per_call = stats.bytes / 10_000
+        @test bytes_per_call == 0.0
+    end
 end