feat: use DI for dense jacobians

Author: avik-pal

src/internal/jacobian.jl

@@ -31,30 +31,25 @@ Construct a cache for the Jacobian of `f` w.r.t. `u`.
 @concrete mutable struct JacobianCache{iip} <: AbstractNonlinearSolveJacobianCache{iip}
     J
     f
-    uf
     fu
     u
     p
-    jac_cache
     alg
     stats::NLStats
     autodiff
     di_extras
+    sdifft_extras
 end
 
 function reinit_cache!(cache::JacobianCache{iip}, args...; p = cache.p,
         u0 = cache.u, kwargs...) where {iip}
     cache.u = u0
     cache.p = p
-    cache.uf = JacobianWrapper{iip}(cache.f, p)
 end
 
 function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
         vjp_autodiff = nothing, jvp_autodiff = nothing, linsolve = missing) where {F}
     iip = isinplace(prob)
-    uf = JacobianWrapper{iip}(f, p)
-
-    autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
 
     has_analytic_jac = SciMLBase.has_jac(f)
     linsolve_needs_jac = concrete_jac(alg) === nothing && (linsolve === missing ||
@@ -65,12 +60,31 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
     @bb fu = similar(fu_)
 
     if !has_analytic_jac && needs_jac
+        autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
         sd = __sparsity_detection_alg(f, autodiff)
-        jac_cache = iip ? sparse_jacobian_cache(autodiff, sd, uf, fu, u) :
-                    sparse_jacobian_cache(
-            autodiff, sd, uf, __maybe_mutable(u, autodiff); fx = fu)
+        sparse_jac = !(sd isa NoSparsityDetection)
+        # Eventually we want to do everything via DI. But for now, we just do the dense via DI
+        if sparse_jac
+            di_extras = nothing
+            uf = JacobianWrapper{iip}(f, p)
+            sdifft_extras = if iip
+                sparse_jacobian_cache(autodiff, sd, uf, fu, u)
+            else
+                sparse_jacobian_cache(
+                    autodiff, sd, uf, __maybe_mutable(u, autodiff); fx = fu)
+            end
+        else
+            sdifft_extras = nothing
+            di_extras = if iip
+                DI.prepare_jacobian(f, fu, autodiff, u, Constant(p))
+            else
+                DI.prepare_jacobian(f, autodiff, u, Constant(p))
+            end
+        end
     else
-        jac_cache = nothing
+        sparse_jac = false
+        di_extras = nothing
+        sdifft_extras = nothing
     end
 
     J = if !needs_jac
@@ -80,36 +94,34 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
             vjp_autodiff, prob, Val(false); check_reverse_mode = false)
         JacobianOperator(prob, fu, u; jvp_autodiff, vjp_autodiff)
     else
-        if has_analytic_jac
-            f.jac_prototype === nothing ?
-            __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u)) :
-            copy(f.jac_prototype)
-        elseif f.jac_prototype === nothing
-            zero(init_jacobian(jac_cache; preserve_immutable = Val(true)))
+        if f.jac_prototype === nothing
+            if !sparse_jac
+                __similar(fu, promote_type(eltype(fu), eltype(u)), length(fu), length(u))
+            else
+                zero(init_jacobian(sdifft_extras; preserve_immutable = Val(true)))
+            end
         else
-            f.jac_prototype
+            similar(f.jac_prototype)
         end
     end
 
     return JacobianCache{iip}(
-        J, f, uf, fu, u, p, jac_cache, alg, stats, autodiff, nothing)
+        J, f, fu, u, p, alg, stats, autodiff, di_extras, sdifft_extras)
 end
 
 function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
         autodiff = nothing, kwargs...) where {F}
     fu = f(u, p)
     if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
-        return JacobianCache{false}(
-            u, f, nothing, fu, u, p, nothing, alg, stats, autodiff, nothing)
+        return JacobianCache{false}(u, f, fu, u, p, alg, stats, autodiff, nothing)
     end
     autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
     di_extras = DI.prepare_derivative(f, autodiff, u, Constant(prob.p))
-    return JacobianCache{false}(
-        u, f, nothing, fu, u, p, nothing, alg, stats, autodiff, di_extras)
+    return JacobianCache{false}(u, f, fu, u, p, alg, stats, autodiff, di_extras, nothing)
 end
 
-@inline (cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p)
-@inline function (cache::JacobianCache)(::Nothing)
+(cache::JacobianCache)(u = cache.u) = cache(cache.J, u, cache.p)
+function (cache::JacobianCache)(::Nothing)
     cache.J isa JacobianOperator &&
         return StatefulJacobianOperator(cache.J, cache.u, cache.p)
     return cache.J
@@ -136,23 +148,31 @@ function (cache::JacobianCache{iip})(
         J::Union{AbstractMatrix, Nothing}, u, p = cache.p) where {iip}
     cache.stats.njacs += 1
     if iip
-        if has_jac(cache.f)
+        if SciMLBase.has_jac(cache.f)
             cache.f.jac(J, u, p)
+        elseif cache.di_extras !== nothing
+            DI.jacobian!(
+                cache.f, cache.fu, J, cache.di_extras, cache.autodiff, u, Constant(p))
         else
-            sparse_jacobian!(J, cache.autodiff, cache.jac_cache, cache.uf, cache.fu, u)
+            uf = JacobianWrapper{iip}(cache.f, p)
+            sparse_jacobian!(J, cache.autodiff, cache.jac_cache, uf, cache.fu, u)
         end
-        J_ = J
+        return J
     else
-        J_ = if has_jac(cache.f)
-            cache.f.jac(u, p)
-        elseif __can_setindex(typeof(J))
-            sparse_jacobian!(J, cache.autodiff, cache.jac_cache, cache.uf, u)
-            J
+        if SciMLBase.has_jac(cache.f)
+            return cache.f.jac(u, p)
+        elseif cache.di_extras !== nothing
+            return DI.jacobian(cache.f, cache.di_extras, cache.autodiff, u, Constant(p))
         else
-            sparse_jacobian(cache.autodiff, cache.jac_cache, cache.uf, u)
+            uf = JacobianWrapper{iip}(cache.f, p)
+            if __can_setindex(typeof(J))
+                sparse_jacobian!(J, cache.autodiff, cache.sdifft_extras, uf, u)
+                return J
+            else
+                return sparse_jacobian(cache.autodiff, cache.sdifft_extras, uf, u)
+            end
         end
     end
-    return J_
 end
 
 # Sparsity Detection Choices
@@ -183,8 +203,7 @@ end
     if SciMLBase.has_colorvec(f)
         return PrecomputedJacobianColorvec(; jac_prototype,
             f.colorvec,
-            partition_by_rows = (ad isa AutoSparse &&
-                                 ADTypes.mode(ad) isa ADTypes.ReverseMode))
+            partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode)
     else
         return JacPrototypeSparsityDetection(; jac_prototype)
     end