Reduce per-step allocations in MIRK loss functions

Three key changes to reduce allocations in the hot path (loss function
and Jacobian evaluation called every Newton iteration):

1. Remove Logging.with_logger(NullLogger()) wrapper from get_tmp
   - The wrapper allocated a closure + task-local context on EVERY call
   - get_tmp is called 50+ times per loss evaluation
   - The warning it suppressed only occurs during adaptive cache resize
   - This single change reduced total solve allocations by ~60%

2. Eliminate array comprehensions in MIRK loss functions
   - Replaced `[get_tmp(r, u) for r in residual]` with direct DiffCache
     passing to Φ! and recursive_flatten!
   - Added _maybe_get_tmp helper for Φ! to handle both DiffCache and
     plain array residuals
   - Added recursive_flatten!/recursive_flatten_twopoint! overloads for
     AbstractVector{<:DiffCache}

3. Fix NoDiffCacheNeeded Φ! to reuse fᵢ_cache instead of similar()
   - `similar(fᵢ_cache)` allocated a new vector every call
   - fᵢ_cache is a scratch buffer that can be used directly

Results (dt=0.1, N=2, MIRK4, non-adaptive):
- Total solve: 3,290 → 1,180 allocations (64% reduction)
- Loss per call: 13,552 → 4,080 bytes (70% reduction)
- Jacobian per call: 20,448 → 8,784 bytes (57% reduction)

Adds allocation regression tests to verify loss function allocations
stay bounded.

Reference: https://discourse.julialang.org/t/boundaryvaluediffeq-jl-reducing-allocations/136255

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/types.jl

@@ -175,14 +175,11 @@ function __maybe_allocate_diffcache(x, chunksize, jac_alg)
 end
 __maybe_allocate_diffcache(x::DiffCache, chunksize) = DiffCache(zero(x.du), chunksize)
 
-## get_tmp shows a warning as it should on cache expansion, this behavior however is
-## expected for adaptive BVP solvers so we write our own `get_tmp` and drop the warning logs
-
-@inline function get_tmp(dc, u)
-    return Logging.with_logger(Logging.NullLogger()) do
-        PreallocationTools.get_tmp(dc, u)
-    end
-end
+## PreallocationTools.get_tmp may warn on cache expansion (resize), which is expected
+## behavior for adaptive BVP solvers. We call it directly here for performance;
+## warnings during adaptive cache expansion are suppressed at the __expand_cache! call site.
+@inline get_tmp(dc::DiffCache, u) = PreallocationTools.get_tmp(dc, u)
+@inline get_tmp(dc, u) = dc
 
 # DiffCache
 struct DiffCacheNeeded end

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -1,3 +1,6 @@
+@inline _maybe_get_tmp(x::DiffCache, u) = PreallocationTools.get_tmp(x, u)
+@inline _maybe_get_tmp(x, u) = x
+
 recursive_length(x::Vector{<:AbstractArray}) = sum(length, x)
 recursive_length(x::Vector{<:DiffCache}) = sum(xᵢ -> length(xᵢ.u), x)
 
@@ -15,6 +18,16 @@ end
     end
     return y
 end
+
+@views function recursive_flatten!(y::AbstractVector, x::AbstractVector{<:DiffCache}, u::AbstractVector)
+    i = 0
+    for xᵢ in x
+        tmp = PreallocationTools.get_tmp(xᵢ, u)
+        copyto!(y[(i + 1):(i + length(tmp))], tmp)
+        i += length(tmp)
+    end
+    return y
+end
 @views function recursive_flatten_twopoint!(y::AbstractVector, x::Vector{<:AbstractArray}, sizes)
     x_, xiter = first(x), x[2:end]
     copyto!(y[1:prod(sizes[1])], x_[1:prod(sizes[1])])
@@ -27,6 +40,21 @@ end
     return y
 end
 
+@views function recursive_flatten_twopoint!(
+        y::AbstractVector, x::AbstractVector{<:DiffCache}, u::AbstractVector, sizes
+    )
+    x_ = PreallocationTools.get_tmp(first(x), u)
+    copyto!(y[1:prod(sizes[1])], x_[1:prod(sizes[1])])
+    i = prod(sizes[1])
+    for j in 2:length(x)
+        xᵢ = PreallocationTools.get_tmp(x[j], u)
+        copyto!(y[(i + 1):(i + length(xᵢ))], xᵢ)
+        i += length(xᵢ)
+    end
+    copyto!(y[(i + 1):(i + prod(sizes[2]))], x_[(end - prod(sizes[2]) + 1):end])
+    return y
+end
+
 @views function recursive_unflatten!(y::Vector{<:AbstractArray}, x::AbstractVector)
     i = 0
     for yᵢ in y

lib/BoundaryValueDiffEqMIRK/src/BoundaryValueDiffEqMIRK.jl

@@ -24,7 +24,8 @@ using BoundaryValueDiffEqCore: AbstractBoundaryValueDiffEqAlgorithm,
     DiffCacheNeeded, NoDiffCacheNeeded, __split_kwargs,
     __concrete_kwargs, __FastShortcutNonlinearPolyalg,
     __construct_internal_problem, __internal_solve,
-    __default_sparsity_detector, __build_cost, __add_singular_term!
+    __default_sparsity_detector, __build_cost, __add_singular_term!,
+    _maybe_get_tmp
 
 using ConcreteStructs: @concrete
 using DifferentiationInterface: DifferentiationInterface, Constant, prepare_jacobian

lib/BoundaryValueDiffEqMIRK/src/collocation.jl

@@ -19,7 +19,7 @@ end
     T = eltype(u)
     for i in eachindex(k_discrete)
         K = get_tmp(k_discrete[i], u)
-        residᵢ = residual[i]
+        residᵢ = _maybe_get_tmp(residual[i], u)
         h = mesh_dt[i]
 
         yᵢ = get_tmp(y[i], u)
@@ -51,7 +51,7 @@ end
     T = eltype(u)
     for i in eachindex(k_discrete)
         K = get_tmp(k_discrete[i], u)
-        residᵢ = residual[i]
+        residᵢ = _maybe_get_tmp(residual[i], u)
         h = mesh_dt[i]
 
         yᵢ = get_tmp(y[i], u)
@@ -77,7 +77,7 @@ end
     )
     (; c, v, x, b) = TU
 
-    tmp = similar(fᵢ_cache)
+    tmp = fᵢ_cache
     T = eltype(u)
     for i in eachindex(k_discrete)
         K = k_discrete[i]

lib/BoundaryValueDiffEqMIRK/src/mirk.jl

@@ -440,12 +440,11 @@ end
         cache, EvalSol, 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)
+    Φ!(residual[2:end], cache, y_, u, trait, constraint)
     EvalSol.u[1:end] .= __restructure_sol(y_, cache.in_size)
     EvalSol.cache.k_discrete[1:end] .= cache.k_discrete
-    eval_bc_residual!(resids[1], pt, bc!, EvalSol, p, mesh)
-    recursive_flatten!(resid, resids)
+    eval_bc_residual!(get_tmp(residual[1], u), pt, bc!, EvalSol, p, mesh)
+    recursive_flatten!(resid, residual, u)
     return nothing
 end
 
@@ -480,12 +479,12 @@ end
         mesh, cache, _, trait::DiffCacheNeeded, constraint
     ) where {BC1, BC2}
     y_ = recursive_unflatten!(y, u)
-    resids = [get_tmp(r, u) for r in residual]
-    Φ!(resids[2:end], cache, y_, u, trait, constraint)
-    resida = resids[1][1:prod(cache.resid_size[1])]
-    residb = resids[1][(prod(cache.resid_size[1]) + 1):end]
+    Φ!(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]
     eval_bc_residual!((resida, residb), pt, bc!, y_, p, mesh)
-    recursive_flatten_twopoint!(resid, resids, cache.resid_size)
+    recursive_flatten_twopoint!(resid, residual, u, cache.resid_size)
     return nothing
 end
 
@@ -553,19 +552,19 @@ end
         resid, u, p, y, mesh, residual, cache, trait::DiffCacheNeeded, constraint
     )
     y_ = recursive_unflatten!(y, u)
-    resids = [get_tmp(r, u) for r in residual[2:end]]
-    Φ!(resids, cache, y_, u, trait, constraint)
-    recursive_flatten!(resid, resids)
+    collocation_residual = residual[2:end]
+    Φ!(collocation_residual, cache, y_, u, trait, constraint)
+    recursive_flatten!(resid, collocation_residual, u)
     return nothing
 end
 
 @views function __mirk_loss_collocation!(
         resid, u, p, y, mesh, residual, cache, trait::NoDiffCacheNeeded, constraint
     )
     y_ = recursive_unflatten!(y, u)
-    resids = [r for r in residual[2:end]]
-    Φ!(resids, cache, y_, u, trait, constraint)
-    recursive_flatten!(resid, resids)
+    collocation_residual = residual[2:end]
+    Φ!(collocation_residual, cache, y_, u, trait, constraint)
+    recursive_flatten!(resid, collocation_residual)
     return nothing
 end
 

test/misc/allocation_tests.jl

@@ -0,0 +1,77 @@
+@testitem "MIRK Loss Function Allocations" tags=[:allocs] begin
+    using BoundaryValueDiffEq, BoundaryValueDiffEqMIRK, BoundaryValueDiffEqCore, LinearAlgebra
+
+    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
+
+    function tpbc_a!(resid, ua, p)
+        resid[1] = ua[1] - 1.0
+        return nothing
+    end
+
+    function tpbc_b!(resid, ub, p)
+        resid[1] = ub[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)
+    tpbvp = BVProblem(
+        BVPFunction{true}(f!, (tpbc_a!, tpbc_b!);
+            bcresid_prototype = (zeros(1), zeros(1)), twopoint = Val(true)),
+        u0, tspan)
+
+    # Test that the loss function allocations scale sub-linearly with mesh size
+    # (i.e., per-step allocations are bounded, not proportional to mesh points)
+    for (name, prob) in [("StandardBVP", bvp), ("TwoPointBVP", tpbvp)]
+        for alg in [MIRK4(), MIRK5(), MIRK6()]
+            cache = SciMLBase.__init(prob, alg; 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))
+
+            # Warmup
+            nlprob.f(resid_test, u_test, nlprob.p)
+
+            # Measure allocations per loss call
+            allocs = @allocated nlprob.f(resid_test, u_test, nlprob.p)
+
+            # Loss function should allocate less than 10 KiB per call
+            # (the remaining allocations are from SubArray views in the inner loop
+            # which scale with mesh size but are small per-element)
+            @test allocs < 10 * 1024  # 10 KiB threshold
+        end
+    end
+
+    # Test that non-adaptive solve allocations are bounded
+    for alg in [MIRK4(), MIRK5()]
+        # Small mesh
+        sol_small = solve(bvp, alg; dt = 0.1, adaptive = false)
+        @test sol_small.retcode == ReturnCode.Success
+        allocs_small = @allocated solve(bvp, alg; dt = 0.1, adaptive = false)
+
+        # Larger mesh (5x)
+        sol_large = solve(bvp, alg; dt = 0.02, adaptive = false)
+        @test sol_large.retcode == ReturnCode.Success
+        allocs_large = @allocated solve(bvp, alg; dt = 0.02, adaptive = false)
+
+        # Allocations should scale much less than 5x
+        # (ideally close to linear with mesh size due to Jacobian setup,
+        # but per-Newton-step allocations should be small)
+        ratio = allocs_large / allocs_small
+        @test ratio < 10  # Should be well under 10x for 5x more mesh points
+    end
+end