Precompile in extensions

Author: avik-pal

ext/SimpleNonlinearSolveDiffResultsForwardDiffExt/SimpleNonlinearSolveDiffResultsForwardDiffExt.jl

@@ -1,5 +1,7 @@
 module SimpleNonlinearSolveDiffResultsForwardDiffExt
 
+import PrecompileTools: @compile_workload, @setup_workload
+
 import ADTypes: AutoForwardDiff, AutoPolyesterForwardDiff
 import ArrayInterface, SciMLBase, SimpleNonlinearSolve, DiffResults, ForwardDiff
 import FastClosures: @closure
@@ -23,4 +25,32 @@ include("jacobian.jl")
 include("hessian.jl")
 include("forward_ad.jl")
 
+@setup_workload begin
+    for T in (Float32, Float64)
+        prob_no_brack_scalar = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))
+        prob_no_brack_iip = NonlinearProblem{true}((du, u, p) -> du .= u .* u .- p,
+            T.([1.0, 1.0, 1.0]), T(2))
+        prob_no_brack_oop = NonlinearProblem{false}((u, p) -> u .* u .- p,
+            T.([1.0, 1.0, 1.0]), T(2))
+
+        algs = [SimpleNonlinearSolve.SimpleNewtonRaphson(; autodiff = AutoForwardDiff()),
+            SimpleNonlinearSolve.SimpleTrustRegion(; autodiff = AutoForwardDiff())]
+
+        algs_no_iip = [SimpleNonlinearSolve.SimpleHalley(; autodiff = AutoForwardDiff())]
+
+        @compile_workload begin
+            for alg in algs
+                solve(prob_no_brack_scalar, alg, abstol = T(1e-2))
+                solve(prob_no_brack_iip, alg, abstol = T(1e-2))
+                solve(prob_no_brack_oop, alg, abstol = T(1e-2))
+            end
+
+            for alg in algs_no_iip
+                solve(prob_no_brack_scalar, alg, abstol = T(1e-2))
+                solve(prob_no_brack_oop, alg, abstol = T(1e-2))
+            end
+        end
+    end
+end
+
 end

ext/SimpleNonlinearSolveDiffResultsForwardDiffExt/jacobian.jl

@@ -92,4 +92,4 @@ function SimpleNonlinearSolve.__value_and_jacobian!(
     end
     out = f(ForwardDiff.Dual{T}(x, one(x)))
     return ForwardDiff.value(out), ForwardDiff.extract_derivative(T, out)
-end
+end

ext/SimpleNonlinearSolveFiniteDiffExt.jl

@@ -1,7 +1,10 @@
 module SimpleNonlinearSolveFiniteDiffExt
 
+import PrecompileTools: @compile_workload, @setup_workload
+
 import ADTypes: AutoFiniteDiff
 import SciMLBase, SimpleNonlinearSolve, FiniteDiff
+import SciMLBase: NonlinearProblem, NonlinearLeastSquaresProblem, solve
 import StaticArraysCore: SArray
 
 @inline SimpleNonlinearSolve.__is_extension_loaded(::Val{:FiniteDiff}) = true
@@ -56,4 +59,32 @@ function SimpleNonlinearSolve.compute_jacobian_and_hessian(
     end
 end
 
+@setup_workload begin
+    for T in (Float32, Float64)
+        prob_no_brack_scalar = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))
+        prob_no_brack_iip = NonlinearProblem{true}((du, u, p) -> du .= u .* u .- p,
+            T.([1.0, 1.0, 1.0]), T(2))
+        prob_no_brack_oop = NonlinearProblem{false}((u, p) -> u .* u .- p,
+            T.([1.0, 1.0, 1.0]), T(2))
+
+        algs = [SimpleNonlinearSolve.SimpleNewtonRaphson(; autodiff = AutoFiniteDiff()),
+            SimpleNonlinearSolve.SimpleTrustRegion(; autodiff = AutoFiniteDiff())]
+
+        algs_no_iip = [SimpleNonlinearSolve.SimpleHalley(; autodiff = AutoFiniteDiff())]
+
+        @compile_workload begin
+            for alg in algs
+                solve(prob_no_brack_scalar, alg, abstol = T(1e-2))
+                solve(prob_no_brack_iip, alg, abstol = T(1e-2))
+                solve(prob_no_brack_oop, alg, abstol = T(1e-2))
+            end
+
+            for alg in algs_no_iip
+                solve(prob_no_brack_scalar, alg, abstol = T(1e-2))
+                solve(prob_no_brack_oop, alg, abstol = T(1e-2))
+            end
+        end
+    end
+end
+
 end

src/SimpleNonlinearSolve.jl

@@ -86,7 +86,6 @@ end
 
 @setup_workload begin
     for T in (Float32, Float64)
-        # FIXME Move this precompilation into the forwarddiff & finitediff extensions
         prob_no_brack_scalar = NonlinearProblem{false}((u, p) -> u .* u .- p, T(0.1), T(2))
         prob_no_brack_iip = NonlinearProblem{true}((du, u, p) -> du .= u .* u .- p,
             T.([1.0, 1.0, 1.0]), T(2))
@@ -96,21 +95,12 @@ end
         algs = [SimpleBroyden(), SimpleKlement(), SimpleDFSane(),
             SimpleLimitedMemoryBroyden(; threshold = 2)]
 
-        # algs = [SimpleNewtonRaphson(), SimpleTrustRegion()]
-
-        # algs_no_iip = [SimpleHalley()]
-
         @compile_workload begin
             for alg in algs
                 solve(prob_no_brack_scalar, alg, abstol = T(1e-2))
                 solve(prob_no_brack_iip, alg, abstol = T(1e-2))
                 solve(prob_no_brack_oop, alg, abstol = T(1e-2))
             end
-
-            # for alg in algs_no_iip
-            #     solve(prob_no_brack_scalar, alg, abstol = T(1e-2))
-            #     solve(prob_no_brack_oop, alg, abstol = T(1e-2))
-            # end
         end
 
         prob_brack = IntervalNonlinearProblem{false}((u, p) -> u * u - p,