Add Polyester ForwardDiff support

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "SimpleNonlinearSolve"
 uuid = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 authors = ["SciML"]
-version = "1.0.4"
+version = "1.1.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -17,8 +17,14 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
+[extensions]
+SimpleNonlinearSolvePolyesterForwardDiffExt = "PolyesterForwardDiff"
+
+[weakdeps]
+PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
+
 [compat]
-ADTypes = "0.2"
+ADTypes = "0.2.6"
 ArrayInterface = "7"
 ConcreteStructs = "0.2"
 DiffEqBase = "6.126"

ext/SimpleNonlinearSolvePolyesterForwardDiffExt.jl

@@ -0,0 +1,13 @@
+module SimpleNonlinearSolvePolyesterForwardDiffExt
+
+using SimpleNonlinearSolve, PolyesterForwardDiff
+
+@inline SimpleNonlinearSolve.__is_extension_loaded(::Val{:PolyesterForwardDiff}) = true
+
+@inline function SimpleNonlinearSolve.__polyester_forwarddiff_jacobian!(f!::F, y, J, x,
+        chunksize) where {F}
+    PolyesterForwardDiff.threaded_jacobian!(f!, y, J, x, chunksize)
+    return J
+end
+
+end

src/SimpleNonlinearSolve.jl

@@ -9,7 +9,7 @@ import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidat
     import DiffEqBase: AbstractNonlinearTerminationMode,
         AbstractSafeNonlinearTerminationMode, AbstractSafeBestNonlinearTerminationMode,
         NonlinearSafeTerminationReturnCode, get_termination_mode,
-        NONLINEARSOLVE_DEFAULT_NORM
+        NONLINEARSOLVE_DEFAULT_NORM, _get_tolerance
     using FiniteDiff, ForwardDiff
     import ForwardDiff: Dual
     import MaybeInplace: @bb, setindex_trait, CanSetindex, CannotSetindex
@@ -23,6 +23,8 @@ abstract type AbstractSimpleNonlinearSolveAlgorithm <: AbstractNonlinearAlgorith
 abstract type AbstractBracketingAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm end
 abstract type AbstractNewtonAlgorithm <: AbstractSimpleNonlinearSolveAlgorithm end
 
+@inline __is_extension_loaded(::Val) = false
+
 include("utils.jl")
 
 ## Nonlinear Solvers

src/nlsolve/halley.jl

@@ -53,14 +53,17 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleHalley, args...;
         fx, dfx, d2fx = compute_jacobian_and_hessian(alg.autodiff, prob, fx, x)
         setindex_trait(x) === CannotSetindex() && (A = dfx)
 
-        aᵢ = dfx \ _vec(fx)
+        # Factorize Once and Reuse
+        dfx_fact = factorize(dfx)
+
+        aᵢ = dfx_fact \ _vec(fx)
         A_ = _vec(A)
         @bb A_ = d2fx × aᵢ
         A = _restructure(A, A_)
 
         @bb Aaᵢ = A × aᵢ
         @bb A .*= -1
-        bᵢ = dfx \ Aaᵢ
+        bᵢ = dfx_fact \ Aaᵢ
 
         cᵢ_ = _vec(cᵢ)
         @bb @. cᵢ_ = (aᵢ * aᵢ) / (-aᵢ + (T(0.5) * bᵢ))

src/nlsolve/raphson.jl

@@ -1,6 +1,6 @@
 """
     SimpleNewtonRaphson(autodiff)
-    SimpleNewtonRaphson(; autodiff = AutoForwardDiff())
+    SimpleNewtonRaphson(; autodiff = nothing)
 
 A low-overhead implementation of Newton-Raphson. This method is non-allocating on scalar
 and static array problems.
@@ -14,10 +14,11 @@ and static array problems.
 ### Keyword Arguments
 
   - `autodiff`: determines the backend used for the Jacobian. Defaults to
-    `AutoForwardDiff()`. Valid choices are `AutoForwardDiff()` or `AutoFiniteDiff()`.
+    `nothing`. Valid choices are `AutoPolyesterForwardDiff()`, `AutoForwardDiff()` or
+    `AutoFiniteDiff()`.
 """
 @kwdef @concrete struct SimpleNewtonRaphson <: AbstractNewtonAlgorithm
-    autodiff = AutoForwardDiff()
+    autodiff = nothing
 end
 
 const SimpleGaussNewton = SimpleNewtonRaphson
@@ -27,14 +28,15 @@ function SciMLBase.__solve(prob::Union{NonlinearProblem, NonlinearLeastSquaresPr
         maxiters = 1000, termination_condition = nothing, alias_u0 = false, kwargs...)
     x = __maybe_unaliased(prob.u0, alias_u0)
     fx = _get_fx(prob, x)
+    autodiff = __get_concrete_autodiff(prob, alg.autodiff)
     @bb xo = copy(x)
-    J, jac_cache = jacobian_cache(alg.autodiff, prob.f, fx, x, prob.p)
+    J, jac_cache = jacobian_cache(autodiff, prob.f, fx, x, prob.p)
 
     abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fx, x,
         termination_condition)
 
     for i in 1:maxiters
-        fx, dfx = value_and_jacobian(alg.autodiff, prob.f, fx, x, prob.p, jac_cache; J)
+        fx, dfx = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J)
 
         if i == 1
             iszero(fx) && build_solution(prob, alg, x, fx; retcode = ReturnCode.Success)

src/nlsolve/trustRegion.jl

@@ -10,7 +10,8 @@ scalar and static array problems.
 ### Keyword Arguments
 
   - `autodiff`: determines the backend used for the Jacobian. Defaults to
-    `AutoForwardDiff()`. Valid choices are `AutoForwardDiff()` or `AutoFiniteDiff()`.
+    `nothing`. Valid choices are `AutoPolyesterForwardDiff()`, `AutoForwardDiff()` or
+    `AutoFiniteDiff()`.
   - `max_trust_radius`: the maximum radius of the trust region. Defaults to
     `max(norm(f(u0)), maximum(u0) - minimum(u0))`.
   - `initial_trust_radius`: the initial trust region radius. Defaults to
@@ -37,7 +38,7 @@ scalar and static array problems.
     row, `max_shrink_times` is exceeded, the algorithm returns. Defaults to `32`.
 """
 @kwdef @concrete struct SimpleTrustRegion <: AbstractNewtonAlgorithm
-    autodiff = AutoForwardDiff()
+    autodiff = nothing
     max_trust_radius = 0.0
     initial_trust_radius = 0.0
     step_threshold = 0.0001
@@ -61,11 +62,12 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleTrustRegion, args.
     t₁ = T(alg.shrink_factor)
     t₂ = T(alg.expand_factor)
     max_shrink_times = alg.max_shrink_times
+    autodiff = __get_concrete_autodiff(prob, alg.autodiff)
 
     fx = _get_fx(prob, x)
     @bb xo = copy(x)
-    J, jac_cache = jacobian_cache(alg.autodiff, prob.f, fx, x, prob.p)
-    fx, ∇f = value_and_jacobian(alg.autodiff, prob.f, fx, x, prob.p, jac_cache; J)
+    J, jac_cache = jacobian_cache(autodiff, prob.f, fx, x, prob.p)
+    fx, ∇f = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J)
 
     abstol, reltol, tc_cache = init_termination_cache(abstol, reltol, fx, x,
         termination_condition)
@@ -116,7 +118,7 @@ function SciMLBase.__solve(prob::NonlinearProblem, alg::SimpleTrustRegion, args.
             # Take the step.
             @bb @. xo = x
 
-            fx, ∇f = value_and_jacobian(alg.autodiff, prob.f, fx, x, prob.p, jac_cache; J)
+            fx, ∇f = value_and_jacobian(autodiff, prob.f, fx, x, prob.p, jac_cache; J)
 
             # Update the trust region radius.
             (r > η₃) && (norm(δ) ≈ Δ) && (Δ = min(t₂ * Δ, Δₘₐₓ))

src/utils.jl

@@ -26,15 +26,6 @@ Return the maximum of `a` and `b` if `x1 > x0`, otherwise return the minimum.
 """
 __max_tdir(a, b, x0, x1) = ifelse(x1 > x0, max(a, b), min(a, b))
 
-__cvt_real(::Type{T}, ::Nothing) where {T} = nothing
-__cvt_real(::Type{T}, x) where {T} = real(T(x))
-
-_get_tolerance(η, ::Type{T}) where {T} = __cvt_real(T, η)
-function _get_tolerance(::Nothing, ::Type{T}) where {T}
-    η = real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)
-    return _get_tolerance(η, T)
-end
-
 __standard_tag(::Nothing, x) = ForwardDiff.Tag(SimpleNonlinearSolveTag(), eltype(x))
 __standard_tag(tag::ForwardDiff.Tag, _) = tag
 __standard_tag(tag, x) = ForwardDiff.Tag(tag, eltype(x))
@@ -60,6 +51,12 @@ function __get_jacobian_config(ad::AutoForwardDiff{CS}, f!, y, x) where {CS}
     return ForwardDiff.JacobianConfig(f!, y, x, ck, tag)
 end
 
+function __get_jacobian_config(ad::AutoPolyesterForwardDiff{CS}, args...) where {CS}
+    x = last(args)
+    return (CS === nothing || CS ≤ 0) ? __pick_forwarddiff_chunk(x) :
+           ForwardDiff.Chunk{CS}()
+end
+
 """
     value_and_jacobian(ad, f, y, x, p, cache; J = nothing)
 
@@ -81,6 +78,9 @@ function value_and_jacobian(ad, f::F, y, x::X, p, cache; J = nothing) where {F,
             FiniteDiff.finite_difference_jacobian!(J, _f, x, cache)
             _f(y, x)
             return y, J
+        elseif ad isa AutoPolyesterForwardDiff
+            __polyester_forwarddiff_jacobian!(_f, y, J, x, cache)
+            return y, J
         else
             throw(ArgumentError("Unsupported AD method: $(ad)"))
         end
@@ -100,12 +100,18 @@ function value_and_jacobian(ad, f::F, y, x::X, p, cache; J = nothing) where {F,
         elseif ad isa AutoFiniteDiff
             J_fd = FiniteDiff.finite_difference_jacobian(_f, x, cache)
             return _f(x), J_fd
+        elseif ad isa AutoPolyesterForwardDiff
+            __polyester_forwarddiff_jacobian!(_f, J, x, cache)
+            return _f(x), J
         else
             throw(ArgumentError("Unsupported AD method: $(ad)"))
         end
     end
 end
 
+# Declare functions
+function __polyester_forwarddiff_jacobian! end
+
 function value_and_jacobian(ad, f::F, y, x::Number, p, cache; J = nothing) where {F}
     if DiffEqBase.has_jac(f)
         return f(x, p), f.jac(x, p)
@@ -132,7 +138,7 @@ function jacobian_cache(ad, f::F, y, x::X, p) where {F, X <: AbstractArray}
         J = similar(y, length(y), length(x))
         if DiffEqBase.has_jac(f)
             return J, nothing
-        elseif ad isa AutoForwardDiff
+        elseif ad isa AutoForwardDiff || ad isa AutoPolyesterForwardDiff
             return J, __get_jacobian_config(ad, _f, y, x)
         elseif ad isa AutoFiniteDiff
             return J, FiniteDiff.JacobianCache(copy(x), copy(y), copy(y), ad.fdtype)
@@ -146,6 +152,10 @@ function jacobian_cache(ad, f::F, y, x::X, p) where {F, X <: AbstractArray}
         elseif ad isa AutoForwardDiff
             J = ArrayInterface.can_setindex(x) ? similar(y, length(y), length(x)) : nothing
             return J, __get_jacobian_config(ad, _f, x)
+        elseif ad isa AutoPolyesterForwardDiff
+            @assert ArrayInterface.can_setindex(x) "PolyesterForwardDiff requires mutable inputs."
+            J = similar(y, length(y), length(x))
+            return J, __get_jacobian_config(ad, _f, x)
         elseif ad isa AutoFiniteDiff
             return nothing, FiniteDiff.JacobianCache(copy(x), copy(y), copy(y), ad.fdtype)
         else
@@ -350,3 +360,18 @@ end
     (alias || !ArrayInterface.can_setindex(typeof(x))) && return x
     return deepcopy(x)
 end
+
+# Decide which AD backend to use
+@inline __get_concrete_autodiff(prob, ad::ADTypes.AbstractADType) = ad
+@inline function __get_concrete_autodiff(prob, ::Nothing)
+    if ForwardDiff.can_dual(eltype(prob.u0))
+        if __is_extension_loaded(Val(:PolyesterForwardDiff)) && !(prob.u0 isa Number) &&
+           ArrayInterface.can_setindex(prob.u0)
+            return AutoPolyesterForwardDiff()
+        else
+            return AutoForwardDiff()
+        end
+    else
+        return AutoFiniteDiff()
+    end
+end