Merge pull request #132 from avik-pal/ap/control_precompile

Author: avik-pal

Project.toml

@@ -14,6 +14,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -46,6 +47,7 @@ ODEInterface = "0.5"
 OrdinaryDiffEq = "6"
 PreallocationTools = "0.4"
 PrecompileTools = "1"
+Preferences = "1"
 RecursiveArrayTools = "2.38.10"
 Reexport = "0.2, 1.0"
 SciMLBase = "2.5"

README.md

@@ -40,3 +40,19 @@ sol = solve(prob, MIRK4(), dt = 0.05)
 ## Available Solvers
 
 For the list of available solvers, please refer to the [DifferentialEquations.jl BVP Solvers page](https://docs.sciml.ai/DiffEqDocs/stable/solvers/bvp_solve/). For options for the `solve` command, see the [common solver options page](https://docs.sciml.ai/DiffEqDocs/stable/basics/common_solver_opts/).
+
+## Controlling Precompilation
+
+Precompilation can be controlled via `Preferences.jl`
+
+- `PrecompileMIRK` -- Precompile the MIRK2 - MIRK6 algorithms (default: `true`).
+- `PrecompileShooting` -- Precompile the single shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
+- `PrecompileMultipleShooting` -- Precompile the multiple shooting algorithms (default: `true`). This is triggered when `OrdinaryDiffEq` is loaded.
+
+To set these preferences before loading the package, do the following (replacing `PrecompileShooting` with the preference you want to set, or pass in multiple pairs to set them together):
+
+```julia
+using Preferences, UUIDs
+Preferences.set_preferences!(UUID("764a87c0-6b3e-53db-9096-fe964310641d"),
+                             "PrecompileShooting" => false)
+```

ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl

@@ -3,6 +3,7 @@ module BoundaryValueDiffEqOrdinaryDiffEqExt
 # This extension doesn't add any new feature atm but is used to precompile some common
 # shooting workflows
 
+import Preferences: @load_preference
 import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidations
 
 @recompile_invalidations begin
@@ -39,16 +40,23 @@ end
         TwoPointBVProblem(f1, (bc1_a, bc1_b), u0, tspan; bcresid_prototype),
     ]
 
-    algs = [
-        Shooting(Tsit5();
-            nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff(chunksize = 2))),
-        MultipleShooting(10,
-            Tsit5();
-            nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff(chunksize = 2)),
-            jac_alg = BVPJacobianAlgorithm(;
-                bc_diffmode = AutoForwardDiff(; chunksize = 2),
-                nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
-    ]
+    algs = []
+
+    if @load_preference("PrecompileShooting", true)
+        push!(algs,
+            Shooting(Tsit5();
+                nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff(chunksize = 2))))
+    end
+
+    if @load_preference("PrecompileMultipleShooting", true)
+        push!(algs,
+            MultipleShooting(10,
+                Tsit5();
+                nlsolve = NewtonRaphson(; autodiff = AutoForwardDiff(chunksize = 2)),
+                jac_alg = BVPJacobianAlgorithm(;
+                    bc_diffmode = AutoForwardDiff(; chunksize = 2),
+                    nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))))
+    end
 
     @compile_workload begin
         for prob in probs, alg in algs

src/BoundaryValueDiffEq.jl

@@ -4,8 +4,8 @@ import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidat
 
 @recompile_invalidations begin
     using ADTypes, Adapt, BandedMatrices, DiffEqBase, ForwardDiff, LinearAlgebra,
-        NonlinearSolve, PreallocationTools, RecursiveArrayTools, Reexport, SciMLBase,
-        Setfield, SparseArrays, SparseDiffTools
+        NonlinearSolve, PreallocationTools, Preferences, RecursiveArrayTools, Reexport,
+        SciMLBase, Setfield, SparseArrays, SparseDiffTools
 
     import ADTypes: AbstractADType
     import ArrayInterface: matrix_colors,
@@ -74,13 +74,18 @@ end
         TwoPointBVProblem(f1, (bc1_a, bc1_b), u0, tspan; bcresid_prototype),
     ]
 
+    algs = []
+
     jac_alg = BVPJacobianAlgorithm(AutoForwardDiff(; chunksize = 2))
 
-    @compile_workload begin
-        for prob in probs,
-            alg in (MIRK2(; jac_alg), MIRK3(; jac_alg), MIRK4(; jac_alg),
-                MIRK5(; jac_alg), MIRK6(; jac_alg))
+    if Preferences.@load_preference("PrecompileMIRK", true)
+        append!(algs,
+            [MIRK2(; jac_alg), MIRK3(; jac_alg), MIRK4(; jac_alg),
+                MIRK5(; jac_alg), MIRK6(; jac_alg)])
+    end
 
+    @compile_workload begin
+        for prob in probs, alg in algs
             solve(prob, alg; dt = 0.2)
         end
     end