Skip to content

Commit

Permalink
fix: only precompile selected workloads for faster loading
Browse files Browse the repository at this point in the history
  • Loading branch information
avik-pal committed Oct 4, 2024
1 parent d87f299 commit a0c3d06
Show file tree
Hide file tree
Showing 2 changed files with 11 additions and 6 deletions.
1 change: 1 addition & 0 deletions lib/SimpleNonlinearSolve/Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ FiniteDiff = "2.24.0"
ForwardDiff = "0.10.36"
InteractiveUtils = "<0.0.1, 1"
LinearAlgebra = "1.10"
LineSearch = "0.1.3"
MaybeInplace = "0.1.4"
NonlinearSolveBase = "1"
PrecompileTools = "1.2"
Expand Down
16 changes: 10 additions & 6 deletions lib/SimpleNonlinearSolve/src/SimpleNonlinearSolve.jl
Original file line number Diff line number Diff line change
Expand Up @@ -75,29 +75,33 @@ end
function solve_adjoint_internal end

@setup_workload begin
for T in (Float32, Float64)
for T in (Float64,)
prob_scalar = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))
prob_iip = NonlinearProblem{true}((du, u, p) -> du .= u .* u .- p, ones(T, 3), T(2))
prob_oop = NonlinearProblem{false}((u, p) -> u .* u .- p, ones(T, 3), T(2))

# Only compile frequently used algorithms -- mostly from the NonlinearSolve default
algs = [
SimpleBroyden(),
# SimpleDFSane(),
SimpleKlement(),
# SimpleLimitedMemoryBroyden(),
SimpleHalley(),
SimpleNewtonRaphson(),
SimpleTrustRegion()
# SimpleHalley(),
SimpleNewtonRaphson(; autodiff = AutoForwardDiff(; chunksize = 1)),
# SimpleTrustRegion()
]

@compile_workload begin
for alg in algs, prob in (prob_scalar, prob_iip, prob_oop)
CommonSolve.solve(prob, alg)
@sync for alg in algs
for prob in (prob_scalar, prob_iip, prob_oop)
Threads.@spawn CommonSolve.solve(prob, alg; abstol = 1e-2)
end
end
end
end
end


export AutoFiniteDiff, AutoForwardDiff, AutoPolyesterForwardDiff

export Alefeld, Bisection, Brent, Falsi, ITP, Ridder
Expand Down

0 comments on commit a0c3d06

Please sign in to comment.