Eliminate dense homotopy step allocations

Cache an owned ForwardDiff destination adapter for dense in-place Jacobians and avoid reshaping already-flat Newton steps. Keep the stored Jacobian matrix identity visible to callers.

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

Author: ChrisRackauckas

lib/NonlinearSolveBase/src/jacobian.jl

@@ -125,7 +125,8 @@ function construct_jacobian_cache(
         end
     end
 
-    return JacobianCache(J, f, fu, p, stats, autodiff, di_extras)
+    ad_buffer = _jacobian_ad_buffer(J, f, fu, u, autodiff)
+    return JacobianCache(J, ad_buffer, f, fu, p, stats, autodiff, di_extras)
 end
 
 function construct_jacobian_cache(
@@ -134,7 +135,7 @@ function construct_jacobian_cache(
         linsolve = missing
     )
     if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
-        return JacobianCache(fu, f, fu, p, stats, autodiff, nothing)
+        return JacobianCache(fu, fu, f, fu, p, stats, autodiff, nothing)
     end
     if autodiff === nothing
         throw(ArgumentError("`autodiff` argument to `construct_jacobian_cache` must be \
@@ -145,11 +146,37 @@ function construct_jacobian_cache(
     @assert !(autodiff isa AutoSparse) "`autodiff` cannot be `AutoSparse` for scalar \
                                         nonlinear problems."
     di_extras = DI.prepare_derivative(f, autodiff, u, Constant(prob.p))
-    return JacobianCache(u, f, fu, p, stats, autodiff, di_extras)
+    return JacobianCache(u, u, f, fu, p, stats, autodiff, di_extras)
+end
+
+# ForwardDiff reshapes its destination before returning it. This adapter lets a dense
+# vector-to-vector Jacobian return the stored matrix directly instead of allocating an
+# equivalent array header on every evaluation.
+struct _JacobianADBuffer{T, M <: Matrix{T}} <: AbstractMatrix{T}
+    data::M
+end
+
+Base.parent(A::_JacobianADBuffer) = A.data
+Base.size(A::_JacobianADBuffer) = size(parent(A))
+Base.axes(A::_JacobianADBuffer) = axes(parent(A))
+Base.IndexStyle(::Type{<:_JacobianADBuffer}) = IndexLinear()
+@inline Base.getindex(A::_JacobianADBuffer, i::Int) = parent(A)[i]
+@inline Base.setindex!(A::_JacobianADBuffer, value, i::Int) = (parent(A)[i] = value)
+function Base.reshape(A::_JacobianADBuffer, dims::Dims)
+    return dims == size(A) ? parent(A) : reshape(parent(A), dims)
+end
+
+_jacobian_ad_buffer(J, f, fu, u, autodiff) = J
+function _jacobian_ad_buffer(
+        J::Matrix, f, fu::Vector, u::Vector, autodiff::AutoForwardDiff
+    )
+    SciMLBase.isinplace(f) && !SciMLBase.has_jac(f) || return J
+    return _JacobianADBuffer(J)
 end
 
 @concrete mutable struct JacobianCache <: AbstractJacobianCache
     J
+    ad_buffer
     f <: NonlinearFunction
     fu
     p
@@ -218,7 +245,8 @@ function (cache::JacobianCache)(u)
             f.jac(J, u, p)
         else
             DI.jacobian!(
-                f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p)
+                f, cache.fu, cache.ad_buffer, cache.di_extras, cache.autodiff, u,
+                Constant(p)
             )
         end
         return J

lib/NonlinearSolveBase/src/utils.jl

@@ -105,6 +105,10 @@ convert_real(::Type{T}, ::Nothing) where {T} = nothing
 convert_real(::Type{T}, x) where {T} = real(T(x))
 
 restructure(::Number, x::Number) = x
+function restructure(y::Vector, x::Vector)
+    size(y) == size(x) && return x
+    return ArrayInterface.restructure(y, x)
+end
 function restructure(
         y::T1, x::T2
     ) where {T1 <: AbstractSciMLOperator, T2 <: AbstractSciMLOperator}

lib/NonlinearSolveFirstOrder/test/sparsity_tests__item2.jl

@@ -28,6 +28,10 @@ for nlf in (f, f!)
 
         cache = init(nlprob, NewtonRaphson(); abstol = 1.0e-9)
         @test cache.jac_cache.J isa Matrix
+        if nlf === f!
+            J = cache.jac_cache.J
+            @test cache.jac_cache(u0) === J
+        end
         sol = solve!(cache)
         @test SciMLBase.successful_retcode(sol)
     end