Track down performance pitfalls and purge dynamic dispatches and allocations

Author: avik-pal

Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0"
 manifest_format = "2.0"
-project_hash = "2fa62d6199f8a6cecd1dab4dd969e2f9c3e4eb5d"
+project_hash = "ee8f38812d75ecf5b51425c9f9559c9e53418c46"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245"
@@ -894,12 +894,6 @@ deps = ["Artifacts", "Libdl", "libblastrampoline_jll"]
 uuid = "bea87d4a-7f5b-5778-9afe-8cc45184846c"
 version = "7.2.1+1"
 
-[[deps.SumTypes]]
-deps = ["MacroTools"]
-git-tree-sha1 = "dc8ae794496a9f04e16393612511223750291547"
-uuid = "8e1ec7a9-0e02-4297-b0fe-6433085c89f2"
-version = "0.5.1"
-
 [[deps.SymbolicIndexingInterface]]
 git-tree-sha1 = "74502f408d99fc217a9d7cd901d9ffe45af892b1"
 uuid = "2efcf032-c050-4f8e-a9bb-153293bab1f5"

Project.toml

@@ -18,6 +18,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 MaybeInplace = "bb5d69b7-63fc-4a16-80bd-7e42200c7bdb"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
@@ -26,7 +27,6 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SparseDiffTools = "47a9eef4-7e08-11e9-0b38-333d64bd3804"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
-SumTypes = "8e1ec7a9-0e02-4297-b0fe-6433085c89f2"
 TimerOutputs = "a759f4b9-e2f1-59dc-863e-4aeb61b1ea8f"
 
 [weakdeps]
@@ -95,7 +95,6 @@ SpeedMapping = "0.3"
 StableRNGs = "1"
 StaticArrays = "1.7"
 StaticArraysCore = "1.4"
-SumTypes = "0.5"
 Sundials = "4.23.1"
 Symbolics = "5.13"
 Test = "1"

docs/src/basics/solve.md

@@ -14,9 +14,13 @@ solve(prob::SciMLBase.NonlinearProblem, args...; kwargs...)
 ## Iteration Controls
 
   - `maxiters::Int`: The maximum number of iterations to perform. Defaults to `1000`.
-  - `maxtime`: The maximum time for solving the nonlinear system of equations. Defaults to `Inf`.
-  - `abstol::Number`: The absolute tolerance. Defaults to `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`.
-  - `reltol::Number`: The relative tolerance. Defaults to `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`.
+  - `maxtime`: The maximum time for solving the nonlinear system of equations. Defaults to
+    `nothing` which means no time limit. Note that setting a time limit does have a small
+    overhead.
+  - `abstol::Number`: The absolute tolerance. Defaults to
+    `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`.
+  - `reltol::Number`: The relative tolerance. Defaults to
+    `real(oneunit(T)) * (eps(real(one(T))))^(4 // 5)`.
   - `termination_condition`: Termination Condition from DiffEqBase. Defaults to
     `AbsSafeBestTerminationMode()` for `NonlinearSolve.jl` and `AbsTerminateMode()` for
     `SimpleNonlinearSolve.jl`.

docs/src/tutorials/code_optimization.md

@@ -115,7 +115,7 @@ to normal array expressions, for example:
 ```@example small_opt
 using StaticArrays
 A = SA[2.0, 3.0, 5.0]
-typeof(A) # SVector{3, Float64} (alias for SArray{Tuple{3}, Float64, 1, 3})
+typeof(A)
 ```
 
 Notice that the `3` after `SVector` gives the size of the `SVector`. It cannot be changed.

docs/src/tutorials/getting_started.md

@@ -194,8 +194,8 @@ solve(prob, GaussNewton(), reltol = 1e-12, abstol = 1e-12)
 
 ## Going Beyond the Basics: How to use the Documentation
 
-Congrats, you now know how to use the basics of NonlinearSolve.jl! However, there is so much more to
-see. Next check out:
+Congrats, you now know how to use the basics of NonlinearSolve.jl! However, there is so much
+more to see. Next check out:
 
   - [Some code optimization tricks to know about with NonlinearSolve.jl](@ref code_optimization)
   - [An iterator interface which lets you step through the solving process step by step](@ref iterator)

src/NonlinearSolve.jl

@@ -9,8 +9,8 @@ import PrecompileTools: @recompile_invalidations, @compile_workload, @setup_work
 
 @recompile_invalidations begin
     using ADTypes, ConcreteStructs, DiffEqBase, FastBroadcast, FastClosures, LazyArrays,
-        LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Printf, SciMLBase,
-        SimpleNonlinearSolve, SparseArrays, SparseDiffTools, SumTypes, TimerOutputs
+        LineSearches, LinearAlgebra, LinearSolve, MaybeInplace, Preferences, Printf,
+        SciMLBase, SimpleNonlinearSolve, SparseArrays, SparseDiffTools, TimerOutputs
 
     import ArrayInterface: undefmatrix, can_setindex, restructure, fast_scalar_indexing
     import DiffEqBase: AbstractNonlinearTerminationMode,
@@ -40,6 +40,7 @@ const True = Val(true)
 const False = Val(false)
 
 include("abstract_types.jl")
+include("timer_outputs.jl")
 include("internal/helpers.jl")
 
 include("descent/newton.jl")

src/abstract_types.jl

@@ -140,11 +140,11 @@ abstract type AbstractNonlinearSolveExtensionAlgorithm <:
               AbstractNonlinearSolveAlgorithm{:Extension} end
 
 """
-    AbstractNonlinearSolveCache{iip}
+    AbstractNonlinearSolveCache{iip, timeit}
 
 Abstract Type for all NonlinearSolve.jl Caches.
 """
-abstract type AbstractNonlinearSolveCache{iip} end
+abstract type AbstractNonlinearSolveCache{iip, timeit} end
 
 SciMLBase.isinplace(::AbstractNonlinearSolveCache{iip}) where {iip} = iip
 

src/core/approximate_jacobian.jl

@@ -46,8 +46,8 @@ end
 
 @inline concrete_jac(::ApproximateJacobianSolveAlgorithm{CJ}) where {CJ} = CJ
 
-@concrete mutable struct ApproximateJacobianSolveCache{INV, GB, iip} <:
-                         AbstractNonlinearSolveCache{iip}
+@concrete mutable struct ApproximateJacobianSolveCache{INV, GB, iip, timeit} <:
+                         AbstractNonlinearSolveCache{iip, timeit}
     # Basic Requirements
     fu
     u
@@ -79,7 +79,7 @@ end
     steps_since_last_reset::Int
 
     # Timer
-    timer::TimerOutput
+    timer
     total_time::Float64   # Simple Counter which works even if TimerOutput is disabled
 
     # Termination & Tracking
@@ -93,8 +93,8 @@ end
 store_inverse_jacobian(::ApproximateJacobianSolveCache{INV}) where {INV} = INV
 
 function __reinit_internal!(cache::ApproximateJacobianSolveCache{INV, GB, iip}, args...;
-        p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf,
-        kwargs...) where {INV, GB, iip}
+        p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000,
+        maxtime = nothing, kwargs...) where {INV, GB, iip}
     if iip
         recursivecopy!(cache.u, u0)
         cache.prob.f(cache.fu, cache.u, p)
@@ -123,12 +123,12 @@ end
 @internal_caches ApproximateJacobianSolveCache :initialization_cache :descent_cache :linesearch_cache :trustregion_cache :update_rule_cache :reinit_rule_cache
 
 function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip},
-        alg::ApproximateJacobianSolveAlgorithm, args...; alias_u0 = false, maxtime = Inf,
-        maxiters = 1000, abstol = nothing, reltol = nothing, linsolve_kwargs = (;),
-        termination_condition = nothing, internalnorm::F = DEFAULT_NORM,
-        kwargs...) where {uType, iip, F}
-    timer = TimerOutput()
-    @timeit_debug timer "cache construction" begin
+        alg::ApproximateJacobianSolveAlgorithm, args...; alias_u0 = false,
+        maxtime = nothing, maxiters = 1000, abstol = nothing, reltol = nothing,
+        linsolve_kwargs = (;), termination_condition = nothing,
+        internalnorm::F = DEFAULT_NORM, kwargs...) where {uType, iip, F}
+    timer = get_timer_output()
+    @static_timeit timer "cache construction" begin
         (; f, u0, p) = prob
         u = __maybe_unaliased(u0, alias_u0)
         fu = evaluate_f(prob, u)
@@ -181,18 +181,18 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip},
         trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du;
             uses_jacobian_inverse = Val(INV), kwargs...)
 
-        return ApproximateJacobianSolveCache{INV, GB, iip}(fu, u, u_cache, p, du, J, alg,
-            prob, initialization_cache, descent_cache, linesearch_cache, trustregion_cache,
-            update_rule_cache, reinit_rule_cache, inv_workspace, 0, 0, 0, alg.max_resets,
-            maxiters, maxtime, alg.max_shrink_times, 0, timer, 0.0, termination_cache,
-            trace, ReturnCode.Default, false, false)
+        return ApproximateJacobianSolveCache{INV, GB, iip, maxtime !== nothing}(fu, u,
+            u_cache, p, du, J, alg, prob, initialization_cache, descent_cache,
+            linesearch_cache, trustregion_cache, update_rule_cache, reinit_rule_cache,
+            inv_workspace, 0, 0, 0, alg.max_resets, maxiters, maxtime, alg.max_shrink_times,
+            0, timer, 0.0, termination_cache, trace, ReturnCode.Default, false, false)
     end
 end
 
 function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip};
         recompute_jacobian::Union{Nothing, Bool} = nothing) where {INV, GB, iip}
     new_jacobian = true
-    @timeit_debug cache.timer "jacobian init/reinit" begin
+    @static_timeit cache.timer "jacobian init/reinit" begin
         if get_nsteps(cache) == 0  # First Step is special ignore kwargs
             J_init = solve!(cache.initialization_cache, cache.fu, cache.u, Val(false))
             if INV
@@ -248,7 +248,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip};
         end
     end
 
-    @timeit_debug cache.timer "descent" begin
+    @static_timeit cache.timer "descent" begin
         if cache.trustregion_cache !== nothing &&
            hasfield(typeof(cache.trustregion_cache), :trust_region)
             δu, descent_success, descent_intermediates = solve!(cache.descent_cache,
@@ -262,19 +262,19 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip};
 
     if descent_success
         if GB === :LineSearch
-            @timeit_debug cache.timer "linesearch" begin
+            @static_timeit cache.timer "linesearch" begin
                 needs_reset, α = solve!(cache.linesearch_cache, cache.u, δu)
             end
             if needs_reset && cache.steps_since_last_reset > 5 # Reset after a burn-in period
                 cache.force_reinit = true
             else
-                @timeit_debug cache.timer "step" begin
+                @static_timeit cache.timer "step" begin
                     @bb axpy!(α, δu, cache.u)
                     evaluate_f!(cache, cache.u, cache.p)
                 end
             end
         elseif GB === :TrustRegion
-            @timeit_debug cache.timer "trustregion" begin
+            @static_timeit cache.timer "trustregion" begin
                 tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu,
                     cache.u, δu, descent_intermediates)
                 if tr_accepted
@@ -289,7 +289,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip};
             end
             α = true
         elseif GB === :None
-            @timeit_debug cache.timer "step" begin
+            @static_timeit cache.timer "step" begin
                 @bb axpy!(1, δu, cache.u)
                 evaluate_f!(cache, cache.u, cache.p)
             end
@@ -313,7 +313,7 @@ function __step!(cache::ApproximateJacobianSolveCache{INV, GB, iip};
         return nothing
     end
 
-    @timeit_debug cache.timer "jacobian update" begin
+    @static_timeit cache.timer "jacobian update" begin
         cache.J = solve!(cache.update_rule_cache, cache.J, cache.fu, cache.u, δu)
         callback_into_cache!(cache)
     end

src/core/generalized_first_order.jl

@@ -49,8 +49,8 @@ end
 
 concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ
 
-@concrete mutable struct GeneralizedFirstOrderAlgorithmCache{iip, GB} <:
-                         AbstractNonlinearSolveCache{iip}
+@concrete mutable struct GeneralizedFirstOrderAlgorithmCache{iip, GB, timeit} <:
+                         AbstractNonlinearSolveCache{iip, timeit}
     # Basic Requirements
     fu
     u
@@ -75,7 +75,7 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ
     max_shrink_times::Int
 
     # Timer
-    timer::TimerOutput
+    timer
     total_time::Float64   # Simple Counter which works even if TimerOutput is disabled
 
     # State Affect
@@ -89,8 +89,8 @@ concrete_jac(::GeneralizedFirstOrderAlgorithm{CJ}) where {CJ} = CJ
 end
 
 function __reinit_internal!(cache::GeneralizedFirstOrderAlgorithmCache{iip}, args...;
-        p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000, maxtime = Inf,
-        kwargs...) where {iip}
+        p = cache.p, u0 = cache.u, alias_u0::Bool = false, maxiters = 1000,
+        maxtime = nothing, kwargs...) where {iip}
     if iip
         recursivecopy!(cache.u, u0)
         cache.prob.f(cache.fu, cache.u, p)
@@ -118,10 +118,11 @@ end
 
 function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip},
         alg::GeneralizedFirstOrderAlgorithm, args...; alias_u0 = false, maxiters = 1000,
-        abstol = nothing, reltol = nothing, maxtime = Inf, termination_condition = nothing,
-        internalnorm = DEFAULT_NORM, linsolve_kwargs = (;), kwargs...) where {uType, iip}
-    timer = TimerOutput()
-    @timeit_debug timer "cache construction" begin
+        abstol = nothing, reltol = nothing, maxtime = nothing,
+        termination_condition = nothing, internalnorm = DEFAULT_NORM, linsolve_kwargs = (;),
+        kwargs...) where {uType, iip}
+    timer = get_timer_output()
+    @static_timeit timer "cache construction" begin
         (; f, u0, p) = prob
         u = __maybe_unaliased(u0, alias_u0)
         fu = evaluate_f(prob, u)
@@ -166,16 +167,16 @@ function SciMLBase.__init(prob::AbstractNonlinearProblem{uType, iip},
 
         trace = init_nonlinearsolve_trace(alg, u, fu, ApplyArray(__zero, J), du; kwargs...)
 
-        return GeneralizedFirstOrderAlgorithmCache{iip, GB}(fu, u, u_cache, p,
-            du, J, alg, prob, jac_cache, descent_cache, linesearch_cache,
+        return GeneralizedFirstOrderAlgorithmCache{iip, GB, maxtime !== nothing}(fu, u,
+            u_cache, p, du, J, alg, prob, jac_cache, descent_cache, linesearch_cache,
             trustregion_cache, 0, 0, maxiters, maxtime, alg.max_shrink_times, timer, 0.0,
             true, termination_cache, trace, ReturnCode.Default, false)
     end
 end
 
 function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB};
         recompute_jacobian::Union{Nothing, Bool} = nothing, kwargs...) where {iip, GB}
-    @timeit_debug cache.timer "jacobian" begin
+    @static_timeit cache.timer "jacobian" begin
         if (recompute_jacobian === nothing || recompute_jacobian) && cache.make_new_jacobian
             J = cache.jac_cache(cache.u)
             new_jacobian = true
@@ -185,7 +186,7 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB};
         end
     end
 
-    @timeit_debug cache.timer "descent" begin
+    @static_timeit cache.timer "descent" begin
         if cache.trustregion_cache !== nothing &&
            hasfield(typeof(cache.trustregion_cache), :trust_region)
             δu, descent_success, descent_intermediates = solve!(cache.descent_cache,
@@ -200,19 +201,19 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB};
     if descent_success
         cache.make_new_jacobian = true
         if GB === :LineSearch
-            @timeit_debug cache.timer "linesearch" begin
+            @static_timeit cache.timer "linesearch" begin
                 linesearch_failed, α = solve!(cache.linesearch_cache, cache.u, δu)
             end
             if linesearch_failed
                 cache.retcode = ReturnCode.InternalLineSearchFailed
                 cache.force_stop = true
             end
-            @timeit_debug cache.timer "step" begin
+            @static_timeit cache.timer "step" begin
                 @bb axpy!(α, δu, cache.u)
                 evaluate_f!(cache, cache.u, cache.p)
             end
         elseif GB === :TrustRegion
-            @timeit_debug cache.timer "trustregion" begin
+            @static_timeit cache.timer "trustregion" begin
                 tr_accepted, u_new, fu_new = solve!(cache.trustregion_cache, J, cache.fu,
                     cache.u, δu, descent_intermediates)
                 if tr_accepted
@@ -230,7 +231,7 @@ function __step!(cache::GeneralizedFirstOrderAlgorithmCache{iip, GB};
                 end
             end
         elseif GB === :None
-            @timeit_debug cache.timer "step" begin
+            @static_timeit cache.timer "step" begin
                 @bb axpy!(1, δu, cache.u)
                 evaluate_f!(cache, cache.u, cache.p)
             end

src/core/generic.jl

@@ -24,7 +24,7 @@ function SciMLBase.solve!(cache::AbstractNonlinearSolveCache)
     update_trace!(cache.trace, get_nsteps(cache), get_u(cache), get_fu(cache), nothing,
         nothing, nothing; last = True)
 
-    stats = SciMLBase.NLStats(get_nf(cache), get_njacs(cache), get_nfactors(cache),
+    stats = ImmutableNLStats(get_nf(cache), get_njacs(cache), get_nfactors(cache),
         get_nsolve(cache), get_nsteps(cache))
 
     return SciMLBase.build_solution(cache.prob, cache.alg, get_u(cache), get_fu(cache);
@@ -45,18 +45,21 @@ Performs one step of the nonlinear solver.
     respectively. For algorithms that don't use jacobian information, this keyword is
     ignored with a one-time warning.
 """
-function SciMLBase.step!(cache::AbstractNonlinearSolveCache, args...; kwargs...)
-    time_start = time()
-    res = @timeit_debug cache.timer "solve" begin
+function SciMLBase.step!(cache::AbstractNonlinearSolveCache{iip, timeit}, args...;
+        kwargs...) where {iip, timeit}
+    timeit && (time_start = time())
+    res = @static_timeit cache.timer "solve" begin
         __step!(cache, args...; kwargs...)
     end
     cache.nsteps += 1
-    cache.total_time += time() - time_start
 
-    if !cache.force_stop && cache.retcode == ReturnCode.Default &&
-       cache.total_time ≥ cache.maxtime
-        cache.retcode = ReturnCode.MaxTime
-        cache.force_stop = true
+    if timeit
+        cache.total_time += time() - time_start
+        if !cache.force_stop && cache.retcode == ReturnCode.Default &&
+           cache.total_time ≥ cache.maxtime
+            cache.retcode = ReturnCode.MaxTime
+            cache.force_stop = true
+        end
     end
 
     return res