Merge pull request #127 from avik-pal/ap/fast_precompile

Try to improve Compilation Speeds

Author: ChrisRackauckas

.github/workflows/CI.yml

@@ -6,6 +6,11 @@ on:
   push:
     branches:
       - master
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 jobs:
   test:
     runs-on: ubuntu-latest

Project.toml

@@ -13,6 +13,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
+PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
@@ -24,31 +25,37 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 
 [weakdeps]
 ODEInterface = "54ca160b-1b9f-5127-a996-1867f4bc2a2c"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 
 [extensions]
 BoundaryValueDiffEqODEInterfaceExt = "ODEInterface"
+BoundaryValueDiffEqOrdinaryDiffEqExt = "OrdinaryDiffEq"
 
 [compat]
 ADTypes = "0.2"
 Adapt = "3"
+Aqua = "0.7"
 ArrayInterface = "7"
 BandedMatrices = "1"
 ConcreteStructs = "0.2"
-DiffEqBase = "6.94.2"
+DiffEqBase = "6.135"
 ForwardDiff = "0.10"
-NonlinearSolve = "2"
+NonlinearSolve = "2.5"
 ODEInterface = "0.5"
+OrdinaryDiffEq = "6"
 PreallocationTools = "0.4"
+PrecompileTools = "1"
 RecursiveArrayTools = "2.38.10"
 Reexport = "0.2, 1.0"
-SciMLBase = "2.2"
+SciMLBase = "2.5"
 Setfield = "1"
-SparseDiffTools = "2.6"
+SparseDiffTools = "2.9"
 TruncatedStacktraces = "1"
 UnPack = "1"
 julia = "1.9"
 
 [extras]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
 ODEInterface = "54ca160b-1b9f-5127-a996-1867f4bc2a2c"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
@@ -58,4 +65,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "SafeTestsets", "ODEInterface"]
+test = ["StaticArrays", "Random", "DiffEqDevTools", "OrdinaryDiffEq", "Test", "SafeTestsets", "ODEInterface", "Aqua"]

README.md

@@ -1,7 +1,12 @@
 # BoundaryValueDiffEq
 
+[![Join the chat at https://julialang.zulipchat.com #sciml-bridged](https://img.shields.io/static/v1?label=Zulip&message=chat&color=9558b2&labelColor=389826)](https://julialang.zulipchat.com/#narrow/stream/279055-sciml-bridged)
+[![Global Docs](https://img.shields.io/badge/docs-SciML-blue.svg)](https://docs.sciml.ai/BoundaryValueDiffEq/stable/)
+
 [![Build Status](https://github.com/SciML/BoundaryValueDiffEq.jl/workflows/CI/badge.svg)](https://github.com/SciML/BoundaryValueDiffEq.jl/actions?query=workflow%3ACI)
 [![codecov](https://codecov.io/gh/SciML/BoundaryValueDiffEq.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/SciML/BoundaryValueDiffEq.jl)
+[![Package Downloads](https://shields.io/endpoint?url=https://pkgs.genieframework.com/api/v1/badge/BoundaryValueDiffEq)](https://pkgs.genieframework.com?packages=BoundaryValueDiffEq)
+[![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
 
 [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor's%20Guide-blueviolet)](https://github.com/SciML/ColPrac)
 [![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle)

ext/BoundaryValueDiffEqODEInterfaceExt.jl

@@ -66,7 +66,7 @@ end
 # BVPSOL
 #-------
 function __solve(prob::BVProblem, alg::BVPSOL; maxiters = 1000, reltol = 1e-3,
-    dt = 0.0, verbose = true, kwargs...)
+        dt = 0.0, verbose = true, kwargs...)
     _test_bvpm2_bvpsol_problem_criteria(prob, prob.problem_type, :BVPSOL)
     @assert isa(prob.p, SciMLBase.NullParameters) "BVPSOL only supports NullParameters!"
     @assert isa(prob.u0, AbstractVector{<:AbstractArray}) "BVPSOL requires a vector of initial guesses!"

ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl

@@ -0,0 +1,60 @@
+module BoundaryValueDiffEqOrdinaryDiffEqExt
+
+# This extension doesn't add any new feature atm but is used to precompile some common
+# shooting workflows
+
+import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidations
+
+@recompile_invalidations begin
+    using BoundaryValueDiffEq, OrdinaryDiffEq
+end
+
+@setup_workload begin
+    function f1!(du, u, p, t)
+        du[1] = u[2]
+        du[2] = 0
+    end
+    f1(u, p, t) = [u[2], 0]
+
+    function bc1!(residual, u, p, t)
+        residual[1] = u(0.0)[1] - 5
+        residual[2] = u(5.0)[1]
+    end
+    bc1(u, p, t) = [u(0.0)[1] - 5, u(5.0)[1]]
+
+    bc1_a!(residual, ua, p) = (residual[1] = ua[1] - 5)
+    bc1_b!(residual, ub, p) = (residual[1] = ub[1])
+
+    bc1_a(ua, p) = [ua[1] - 5]
+    bc1_b(ub, p) = [ub[1]]
+
+    tspan = (0.0, 5.0)
+    u0 = [5.0, -3.5]
+    bcresid_prototype = (Array{Float64}(undef, 1), Array{Float64}(undef, 1))
+
+    probs = [
+        BVProblem(f1!, bc1!, u0, tspan),
+        BVProblem(f1, bc1, u0, tspan),
+        TwoPointBVProblem(f1!, (bc1_a!, bc1_b!), u0, tspan; bcresid_prototype),
+        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))),
+    ]
+
+    @compile_workload begin
+        for prob in probs, alg in algs
+            solve(prob, alg)
+        end
+    end
+end
+
+end

src/BoundaryValueDiffEq.jl

@@ -1,19 +1,26 @@
 module BoundaryValueDiffEq
 
-using Adapt, BandedMatrices, ForwardDiff, LinearAlgebra, PreallocationTools,
-    RecursiveArrayTools, Reexport, Setfield, SparseArrays
-@reexport using ADTypes, DiffEqBase, NonlinearSolve, SparseDiffTools, SciMLBase
+import PrecompileTools: @compile_workload, @setup_workload, @recompile_invalidations
+
+@recompile_invalidations begin
+    using ADTypes, Adapt, BandedMatrices, DiffEqBase, ForwardDiff, LinearAlgebra,
+        NonlinearSolve, PreallocationTools, RecursiveArrayTools, Reexport, SciMLBase,
+        Setfield, SparseArrays, SparseDiffTools
+
+    import ADTypes: AbstractADType
+    import ArrayInterface: matrix_colors,
+        parameterless_type, undefmatrix, fast_scalar_indexing
+    import ConcreteStructs: @concrete
+    import DiffEqBase: solve
+    import ForwardDiff: pickchunksize
+    import RecursiveArrayTools: ArrayPartition, DiffEqArray
+    import SciMLBase: AbstractDiffEqInterpolation, StandardBVProblem, __solve, _unwrap_val
+    import SparseDiffTools: AbstractSparseADType
+    import TruncatedStacktraces: @truncate_stacktrace
+    import UnPack: @unpack
+end
 
-import ADTypes: AbstractADType
-import ArrayInterface: matrix_colors, parameterless_type, undefmatrix, fast_scalar_indexing
-import ConcreteStructs: @concrete
-import DiffEqBase: solve
-import ForwardDiff: pickchunksize
-import RecursiveArrayTools: ArrayPartition, DiffEqArray
-import SciMLBase: AbstractDiffEqInterpolation, StandardBVProblem, __solve, _unwrap_val
-import SparseDiffTools: AbstractSparseADType
-import TruncatedStacktraces: @truncate_stacktrace
-import UnPack: @unpack
+@reexport using ADTypes, DiffEqBase, NonlinearSolve, SparseDiffTools, SciMLBase
 
 include("types.jl")
 include("utils.jl")
@@ -37,6 +44,48 @@ function __solve(prob::BVProblem, alg::BoundaryValueDiffEqAlgorithm, args...; kw
     return solve!(cache)
 end
 
+@setup_workload begin
+    function f1!(du, u, p, t)
+        du[1] = u[2]
+        du[2] = 0
+    end
+    f1(u, p, t) = [u[2], 0]
+
+    function bc1!(residual, u, p, t)
+        residual[1] = u[1][1] - 5
+        residual[2] = u[end][1]
+    end
+    bc1(u, p, t) = [u[1][1] - 5, u[end][1]]
+
+    bc1_a!(residual, ua, p) = (residual[1] = ua[1] - 5)
+    bc1_b!(residual, ub, p) = (residual[1] = ub[1])
+
+    bc1_a(ua, p) = [ua[1] - 5]
+    bc1_b(ub, p) = [ub[1]]
+
+    tspan = (0.0, 5.0)
+    u0 = [5.0, -3.5]
+    bcresid_prototype = (Array{Float64}(undef, 1), Array{Float64}(undef, 1))
+
+    probs = [
+        BVProblem(f1!, bc1!, u0, tspan),
+        BVProblem(f1, bc1, u0, tspan),
+        TwoPointBVProblem(f1!, (bc1_a!, bc1_b!), u0, tspan; bcresid_prototype),
+        TwoPointBVProblem(f1, (bc1_a, bc1_b), u0, tspan; bcresid_prototype),
+    ]
+
+    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))
+
+            solve(prob, alg; dt = 0.2)
+        end
+    end
+end
+
 export Shooting, MultipleShooting
 export MIRK2, MIRK3, MIRK4, MIRK5, MIRK6
 export MIRKJacobianComputationAlgorithm, BVPJacobianAlgorithm

src/adaptivity.jl

@@ -84,7 +84,7 @@ end
 Generate a new mesh based on the `ŝ`.
 """
 function redistribute!(cache::MIRKCache{iip, T}, Nsub_star, ŝ, mesh,
-    mesh_dt) where {iip, T}
+        mesh_dt) where {iip, T}
     N = length(mesh)
     ζ = sum(ŝ .* mesh_dt) / Nsub_star
     k, i = 1, 0
@@ -229,7 +229,7 @@ function sum_stages!(cache::MIRKCache, w, w′, i::Int, dt = cache.mesh_dt[i])
     sum_stages!(cache.fᵢ_cache.du, cache.fᵢ₂_cache, cache, w, w′, i, dt)
 end
 
-function sum_stages!(z, cache::MIRKCache, w, i::Int, dt = cache.mesh_dt[i])
+function sum_stages!(z::AbstractArray, cache::MIRKCache, w, i::Int, dt = cache.mesh_dt[i])
     @unpack M, stage, mesh, k_discrete, k_interp, mesh_dt = cache
     @unpack s_star = cache.ITU
 

src/algorithms.jl

@@ -41,7 +41,7 @@ function update_nshoots(alg::MultipleShooting, nshoots::Int)
 end
 
 function MultipleShooting(nshoots::Int, ode_alg; nlsolve = NewtonRaphson(),
-    grid_coarsening = true, jac_alg = BVPJacobianAlgorithm())
+        grid_coarsening = true, jac_alg = BVPJacobianAlgorithm())
     @assert grid_coarsening isa Bool || grid_coarsening isa Function ||
             grid_coarsening isa AbstractVector{<:Integer} ||
             grid_coarsening isa NTuple{N, <:Integer} where {N}

src/collocation.jl

@@ -4,7 +4,7 @@ function Φ!(residual, cache::MIRKCache, y, u, p = cache.p)
 end
 
 @views function Φ!(residual, fᵢ_cache, k_discrete, f!, TU::MIRKTableau, y, u, p,
-    mesh, mesh_dt, stage::Int)
+        mesh, mesh_dt, stage::Int)
     @unpack c, v, x, b = TU
 
     tmp = get_tmp(fᵢ_cache, u)
@@ -35,7 +35,7 @@ function Φ(cache::MIRKCache, y, u, p = cache.p)
 end
 
 @views function Φ(fᵢ_cache, k_discrete, f, TU::MIRKTableau, y, u, p, mesh, mesh_dt,
-    stage::Int)
+        stage::Int)
     @unpack c, v, x, b = TU
     residuals = [similar(yᵢ) for yᵢ in y[1:(end - 1)]]
     tmp = get_tmp(fᵢ_cache, u)

src/interpolation.jl

@@ -19,7 +19,7 @@ end
 # FIXME: Fix the interpolation outside the tspan
 
 @inline function interpolation(tvals, id::I, idxs, deriv::D, p,
-    continuity::Symbol = :left) where {I, D}
+        continuity::Symbol = :left) where {I, D}
     @unpack t, u, cache = id
     tdir = sign(t[end] - t[1])
     idx = sortperm(tvals, rev = tdir < 0)
@@ -41,7 +41,7 @@ end
 end
 
 @inline function interpolation!(vals, tvals, id::I, idxs, deriv::D, p,
-    continuity::Symbol = :left) where {I, D}
+        continuity::Symbol = :left) where {I, D}
     @unpack t, cache = id
     tdir = sign(t[end] - t[1])
     idx = sortperm(tvals, rev = tdir < 0)
@@ -54,7 +54,7 @@ end
 end
 
 @inline function interpolation(tval::Number, id::I, idxs, deriv::D, p,
-    continuity::Symbol = :left) where {I, D}
+        continuity::Symbol = :left) where {I, D}
     z = similar(id.cache.fᵢ₂_cache)
     interp_eval!(z, id.cache, tval, id.cache.mesh, id.cache.mesh_dt)
     return z