NonlinearSolveAlg performance testing and improvements

[ChrisRackauckas]

Test cases:

using NonlinearSolve, MINPACK, ModelingToolkit, OrdinaryDiffEqBDF, OrdinaryDiffEqNonlinearSolve, Sundials, LSODA
using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
using BenchmarkTools
using ModelingToolkit: t_nounits as t, D_nounits as D
RUS = RadiusUpdateSchemes

@parameters k₁ k₂ k₃
@variables y₁(t) y₂(t) y₃(t)

eqs = [D(y₁) ~ -k₁ * y₁ + k₃ * y₂ * y₃,
    D(y₂) ~ k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃,
    D(y₃) ~ k₂ * y₂^2]
rober = ODESystem(eqs, t; name = :rober) |> structural_simplify |> complete
prob = ODEProblem(rober, [y₁, y₂, y₃] .=> [1.0; 0.0; 0.0], (0.0, 1e5), [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4), jac = true)
nlalg = NonlinearSolveAlg(NewtonRaphson(autodiff = AutoFiniteDiff()))
nlalgrobust = NonlinearSolveAlg(TrustRegion(autodiff = AutoFiniteDiff()))
sol = solve(prob, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg));

@btime solve(prob, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg));
@btime sol = solve(prob, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg));
@btime sol = solve(prob, FBDF(autodiff=AutoFiniteDiff()));
@btime sol = solve(prob, CVODE_BDF());

@profview for i in 1:100 solve(prob, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalgrobust)) end

integ = init(prob, FBDF(autodiff=AutoFiniteDiff(), nlsolve = NonlinearSolveAlg(NewtonRaphson(autodiff = AutoFiniteDiff()))))
(;ts, u_history, order, u_corrector, bdf_coeffs, r, nlsolver , weights, terk_tmp, terkp1_tmp, atmp, tmp, equi_ts, u₀, ts_tmp, step_limiter!) = integ.cache;
(;z, tmp, ztmp, γ, α, cache, method) = nlsolver;

@report_opt DiffEqBase.__init(cache.prob, nlsolver.alg.alg)
using JET
@test_opt NonlinearSolveBase.select_jacobian_autodiff(cache.prob, nlsolver.alg.alg.autodiff)
alg = NonlinearSolveBase.select_jacobian_autodiff(cache.prob, nlsolver.alg.alg.autodiff)
@test_opt alg.jvp_autodiff

@test_opt NonlinearSolveBase.select_forward_mode_autodiff(cache.prob, nlsolver.alg.alg.autodiff)
@test_opt NonlinearSolveBase.select_reverse_mode_autodiff(cache.prob, nlsolver.alg.alg.vjp_autodiff)

@which NonlinearSolveBase.select_reverse_mode_autodiff(cache.prob, nlsolver.alg.alg.vjp_autodiff)
@which OrdinaryDiffEqNonlinearSolve.step!(nlsolver.cache.cache)
@which NonlinearSolveBase.InternalAPI.step!(nlsolver.cache.cache)
@profview for i in 1:100 OrdinaryDiffEqBDF.nlsolve!(nlsolver, integ, integ.cache, false) end



@descend NonlinearSolveBase.InternalAPI.solve!(
                cache.descent_cache, J, cache.fu, cache.u; new_jacobian, cache.kwargs...
            )


using LinearAlgebra, SparseArrays

const N = 32
const xyd_brusselator = range(0, stop = 1, length = N)
brusselator_f(x, y, t) = (((x - 0.3)^2 + (y - 0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.0
limit(a, N) = a == N + 1 ? 1 : a == 0 ? N : a
function brusselator_2d_loop(du, u, p, t)
    A, B, alpha, dx = p
    alpha = alpha / dx^2
    @inbounds for I in CartesianIndices((N, N))
        i, j = Tuple(I)
        x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
        ip1, im1, jp1, jm1 = limit(i + 1, N), limit(i - 1, N), limit(j + 1, N),
        limit(j - 1, N)
        du[i, j, 1] = alpha * (u[im1, j, 1] + u[ip1, j, 1] + u[i, jp1, 1] + u[i, jm1, 1] -
                       4u[i, j, 1]) +
                      B + u[i, j, 1]^2 * u[i, j, 2] - (A + 1) * u[i, j, 1] +
                      brusselator_f(x, y, t)
        du[i, j, 2] = alpha * (u[im1, j, 2] + u[ip1, j, 2] + u[i, jp1, 2] + u[i, jm1, 2] -
                       4u[i, j, 2]) +
                      A * u[i, j, 1] - u[i, j, 1]^2 * u[i, j, 2]
    end
end
p = (3.4, 1.0, 10.0, step(xyd_brusselator))

function init_brusselator_2d(xyd)
    N = length(xyd)
    u = zeros(N, N, 2)
    for I in CartesianIndices((N, N))
        x = xyd[I[1]]
        y = xyd[I[2]]
        u[I, 1] = 22 * (y * (1 - y))^(3 / 2)
        u[I, 2] = 27 * (x * (1 - x))^(3 / 2)
    end
    u
end
u0 = init_brusselator_2d(xyd_brusselator)
prob2 = ODEProblem(brusselator_2d_loop, u0, (0.0, 11.5), p)

sol = solve(prob2, FBDF(autodiff=false, nlsolve = nlalgs[1]), dt = 1e-6);
sol = solve(prob2, FBDF(autodiff=false));

@btime sol = solve(prob2, FBDF(autodiff=false, nlsolve = nlalgs[3]), dt = 1e-6);
@btime sol = solve(prob2, FBDF(autodiff=false));
@btime sol = solve(prob2, FBDF());
@btime sol = solve(prob2, FBDF(autodiff=false, nlsolve = nlalgs[3]));

@btime sol = solve(prob2, FBDF(autodiff=false, nlsolve = nlalgs[1]), dt=1e-6);

[ChrisRackauckas]

One major remaining piece is that ztmp .= nlstep_data.nlprobmap(nlstepsol) is out of place. @AayushSabharwal can you look into getting an inplace function for that?

[ChrisRackauckas]

My current test case:

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
using NonlinearSolve, OrdinaryDiffEqBDF, OrdinaryDiffEqSDIRK, DiffEqDevTools
using OrdinaryDiffEqNonlinearSolve: NonlinearSolveAlg
using Test
using BenchmarkTools

@parameters k₁ k₂ k₃
@variables y₁(t) y₂(t) y₃(t)

eqs = [D(y₁) ~ -k₁ * y₁ + k₃ * y₂ * y₃,
    D(y₂) ~ k₁ * y₁ - k₂ * y₂^2 - k₃ * y₂ * y₃,
    D(y₃) ~ k₂ * y₂^2]
@mtkcompile rober = ODESystem(eqs, t)
prob = ODEProblem(rober, [[y₁, y₂, y₃] .=> [1.0; 0.0; 0.0]; [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4)], (0.0, 1e5), jac = true)
prob2 = ODEProblem(rober, [[y₁, y₂, y₃] .=> [1.0; 0.0; 0.0]; [k₁, k₂, k₃] .=> (0.04, 3e7, 1e4)], (0.0, 1e5), jac = true, nlstep = true)

nlalg = NonlinearSolveAlg(NewtonRaphson(autodiff = AutoFiniteDiff()));
nlalgrobust = NonlinearSolveAlg(TrustRegion(autodiff = AutoFiniteDiff()));
sol2 = solve(prob2, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg));

@profview for i in 1:10000 solve(prob2, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg); save_everystep=false); end

using Plots
plot(sol)

@btime solve(prob, FBDF(autodiff=AutoFiniteDiff()); save_everystep=false);
@btime solve(prob, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg); save_everystep=false);
@btime solve(prob2, FBDF(autodiff=AutoFiniteDiff(), nlsolve = nlalg); save_everystep=false);

[ChrisRackauckas]

The other loop allocations are fixed in SciML/NonlinearSolve.jl#640 and #2760