feat: use DI for dense jacobians

avik-pal · avik-pal · commit a8e60d5f379d · 2024-09-27T16:31:53.000-04:00
fix: stop storing extra stuff in JacobianCache

feat: use DI for dense jacobians

refactor: remove alg from the cache

fix: don't ignore sparsity for dense_ad
diff --git a/Project.toml b/Project.toml
@@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ConcreteStructs = "2569d6c7-a4a2-43d3-a901-331e8e4be471"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
+DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 FastClosures = "9aa1b823-49e4-5ca5-8b0f-3971ec8bab6a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
@@ -65,8 +66,9 @@ BandedMatrices = "1.5"
 BenchmarkTools = "1.4"
 CUDA = "5.2"
 ConcreteStructs = "0.2.3"
+DifferentiationInterface = "0.6.1"
 DiffEqBase = "6.149.0"
-Enzyme = "0.12"
+Enzyme = "0.12, 0.13"
 ExplicitImports = "1.5"
 FastBroadcast = "0.2.8, 0.3"
 FastClosures = "0.3.2"
diff --git a/src/NonlinearSolve.jl b/src/NonlinearSolve.jl
@@ -7,8 +7,8 @@ end
 using Reexport: @reexport
 using PrecompileTools: @compile_workload, @setup_workload
 
-using ADTypes: ADTypes, AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff,
-               AutoZygote, AutoEnzyme, AutoSparse
+using ADTypes: ADTypes, AbstractADType, AutoFiniteDiff, AutoForwardDiff,
+               AutoPolyesterForwardDiff, AutoZygote, AutoEnzyme, AutoSparse
 # FIXME: deprecated, remove in future
 using ADTypes: AutoSparseFiniteDiff, AutoSparseForwardDiff, AutoSparsePolyesterForwardDiff,
                AutoSparseZygote
@@ -22,6 +22,7 @@ using DiffEqBase: DiffEqBase, AbstractNonlinearTerminationMode,
                   NormTerminationMode, RelNormTerminationMode, RelSafeBestTerminationMode,
                   RelSafeTerminationMode, RelTerminationMode,
                   SimpleNonlinearSolveTerminationMode, SteadyStateDiffEqTerminationMode
+using DifferentiationInterface: DifferentiationInterface, Constant
 using FastBroadcast: @..
 using FastClosures: @closure
 using FiniteDiff: FiniteDiff
@@ -39,7 +40,7 @@ using Printf: @printf
 using Preferences: Preferences, @load_preference, @set_preferences!
 using RecursiveArrayTools: recursivecopy!, recursivefill!
 using SciMLBase: AbstractNonlinearAlgorithm, JacobianWrapper, AbstractNonlinearProblem,
-                 AbstractSciMLOperator, _unwrap_val, has_jac, isinplace, NLStats
+                 AbstractSciMLOperator, _unwrap_val, isinplace, NLStats
 using SciMLJacobianOperators: AbstractJacobianOperator, JacobianOperator, VecJacOperator,
                               JacVecOperator, StatefulJacobianOperator
 using SparseArrays: AbstractSparseMatrix, SparseMatrixCSC
@@ -55,6 +56,8 @@ using SymbolicIndexingInterface: SymbolicIndexingInterface, ParameterIndexingPro
 
 @reexport using SciMLBase, SimpleNonlinearSolve
 
+const DI = DifferentiationInterface
+
 # Type-Inference Friendly Check for Extension Loading
 is_extension_loaded(::Val) = false
 
diff --git a/src/internal/jacobian.jl b/src/internal/jacobian.jl
@@ -31,35 +31,24 @@ 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
-    vjp_autodiff
-    jvp_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)
-    jvp_autodiff = get_concrete_forward_ad(
-        jvp_autodiff, prob, Val(false); check_forward_mode = true)
-    vjp_autodiff = get_concrete_reverse_ad(
-        vjp_autodiff, prob, Val(false); check_reverse_mode = false)
 
     has_analytic_jac = SciMLBase.has_jac(f)
     linsolve_needs_jac = concrete_jac(alg) === nothing && (linsolve === missing ||
@@ -70,90 +59,128 @@ function JacobianCache(prob, alg, f::F, fu_, u, p; stats, autodiff = nothing,
     @bb fu = similar(fu_)
 
     if !has_analytic_jac && needs_jac
-        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)
+        autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
+        sd = sparsity_detection_alg(f, autodiff)
+        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
+        jvp_autodiff = get_concrete_forward_ad(
+            jvp_autodiff, prob, Val(false); check_forward_mode = true)
+        vjp_autodiff = get_concrete_reverse_ad(
+            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, vjp_autodiff, jvp_autodiff)
+        J, f, fu, u, p, stats, autodiff, di_extras, sdifft_extras)
 end
 
 function JacobianCache(prob, alg, f::F, ::Number, u::Number, p; stats,
         autodiff = nothing, kwargs...) where {F}
-    uf = JacobianWrapper{false}(f, p)
-    autodiff = get_concrete_forward_ad(autodiff, prob; check_forward_mode = false)
-    if !(autodiff isa AutoForwardDiff ||
-         autodiff isa AutoPolyesterForwardDiff ||
-         autodiff isa AutoFiniteDiff)
-        # Other cases are not properly supported so we fallback to finite differencing
-        @warn "Scalar AD is supported only for AutoForwardDiff and AutoFiniteDiff. \
-               Detected $(autodiff). Falling back to AutoFiniteDiff."
-        autodiff = AutoFiniteDiff()
+    fu = f(u, p)
+    if SciMLBase.has_jac(f) || SciMLBase.has_vjp(f) || SciMLBase.has_jvp(f)
+        return JacobianCache{false}(u, f, fu, u, p, stats, autodiff, nothing)
     end
-    return JacobianCache{false}(
-        u, f, uf, u, u, p, nothing, alg, stats, autodiff, nothing, nothing)
+    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, fu, u, p, 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
 end
 
+# Operator
 function (cache::JacobianCache)(J::JacobianOperator, u, p = cache.p)
     return StatefulJacobianOperator(J, u, p)
 end
+# Numbers
 function (cache::JacobianCache)(::Number, u, p = cache.p) # Scalar
     cache.stats.njacs += 1
-    J = last(__value_derivative(cache.autodiff, cache.uf, u))
-    return J
+    if SciMLBase.has_jac(cache.f)
+        return cache.f.jac(u, p)
+    elseif SciMLBase.has_vjp(cache.f)
+        return cache.f.vjp(one(u), u, p)
+    elseif SciMLBase.has_jvp(cache.f)
+        return cache.f.jvp(one(u), u, p)
+    end
+    return DI.derivative(cache.f, cache.di_extras, cache.autodiff, u, Constant(p))
 end
-# Compute the Jacobian
+# Actually Compute the Jacobian
 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
-@inline __sparsity_detection_alg(_, _) = NoSparsityDetection()
-@inline function __sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
+function sparsity_detection_alg(f::NonlinearFunction, ad::AbstractADType)
+    # TODO: Also handle case where colorvec is provided
+    f.sparsity === nothing && return NoSparsityDetection()
+    return sparsity_detection_alg(f, AutoSparse(ad; sparsity_detector = f.sparsity))
+end
+
+function sparsity_detection_alg(f::NonlinearFunction, ad::AutoSparse)
     if f.sparsity === nothing
         if f.jac_prototype === nothing
             is_extension_loaded(Val(:Symbolics)) && return SymbolicsSparsityDetection()
@@ -177,28 +204,9 @@ end
     end
 
     if SciMLBase.has_colorvec(f)
-        return PrecomputedJacobianColorvec(; jac_prototype,
-            f.colorvec,
-            partition_by_rows = (ad isa AutoSparse &&
-                                 ADTypes.mode(ad) isa ADTypes.ReverseMode))
+        return PrecomputedJacobianColorvec(; jac_prototype, f.colorvec,
+            partition_by_rows = ADTypes.mode(ad) isa ADTypes.ReverseMode)
     else
         return JacPrototypeSparsityDetection(; jac_prototype)
     end
 end
-
-@inline function __value_derivative(
-        ::Union{AutoForwardDiff, AutoPolyesterForwardDiff}, f::F, x::R) where {F, R}
-    T = typeof(ForwardDiff.Tag(f, R))
-    out = f(ForwardDiff.Dual{T}(x, one(x)))
-    return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
-end
-
-@inline function __value_derivative(ad::AutoFiniteDiff, f::F, x::R) where {F, R}
-    return f(x), FiniteDiff.finite_difference_derivative(f, x, ad.fdtype)
-end
-
-@inline function __scalar_jacvec(f::F, x::R, v::V) where {F, R, V}
-    T = typeof(ForwardDiff.Tag(f, R))
-    out = f(ForwardDiff.Dual{T}(x, v))
-    return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
-end
