A bit less runtime dispatches

Author: avik-pal

Manifest.toml

@@ -802,7 +802,7 @@ version = "1.10.0"
 
 [[deps.SparseDiffTools]]
 deps = ["ADTypes", "Adapt", "ArrayInterface", "Compat", "DataStructures", "FiniteDiff", "ForwardDiff", "Graphs", "LinearAlgebra", "PackageExtensionCompat", "Random", "Reexport", "SciMLOperators", "Setfield", "SparseArrays", "StaticArrayInterface", "StaticArrays", "Tricks", "UnPack", "VertexSafeGraphs"]
-git-tree-sha1 = "766a1f0d19232243de203b84d1a7713f6624a1e4"
+git-tree-sha1 = "fcc87cf2750372313d0f761b7dbd75b559c646a4"
 repo-rev = "patch-1"
 repo-url = "https://github.com/avik-pal/SparseDiffTools.jl"
 uuid = "47a9eef4-7e08-11e9-0b38-333d64bd3804"

ext/BoundaryValueDiffEqOrdinaryDiffEqExt.jl

@@ -39,8 +39,17 @@ end
         TwoPointBVProblem(f1, (bc1_a, bc1_b), u0, tspan; bcresid_prototype),
     ]
 
+    algs = [
+        Shooting(Tsit5()),
+        MultipleShooting(10,
+            Tsit5();
+            jac_alg = BVPJacobianAlgorithm(;
+                bc_diffmode = AutoForwardDiff(; chunksize = 2),
+                nonbc_diffmode = AutoSparseForwardDiff(; chunksize = 2))),
+    ]
+
     @compile_workload begin
-        for prob in probs, alg in (Shooting(Tsit5()), MultipleShooting(10, Tsit5()))
+        for prob in probs, alg in algs
             solve(prob, alg)
         end
     end

src/solve/multiple_shooting.jl

@@ -68,13 +68,8 @@ function __solve_nlproblem!(alg::MultipleShooting, bcresid_prototype, u_at_nodes
     resid_prototype = vcat(bcresid_prototype[1],
         similar(u_at_nodes, cur_nshoot * N), bcresid_prototype[2])
 
-    __resid_nodes = resid_prototype[(resida_len + 1):(resida_len + cur_nshoot * N)]
-    resid_nodes = __maybe_allocate_diffcache(__resid_nodes,
-        pickchunksize((cur_nshoot + 1) * N), alg.jac_alg.diffmode)
-
     loss_fn = (du, u, p) -> __multiple_shooting_2point_loss!(du, u, p, cur_nshoot,
-        nodes, iip, solve_internal_odes!, resida_len, residb_len, N, bca,
-        bcb)
+        nodes, iip, solve_internal_odes!, resida_len, residb_len, N, bca, bcb)
     loss_fnₚ = (du, u) -> loss_fn(du, u, prob.p)
 
     sd_bvp = alg.jac_alg.diffmode isa AbstractSparseADType ?
@@ -113,7 +108,7 @@ function __solve_nlproblem!(alg::MultipleShooting, bcresid_prototype, u_at_nodes
 
     loss_fn = (du, u, p) -> __multiple_shooting_mpoint_loss!(du, u, p, cur_nshoot,
         nodes, iip, solve_internal_odes!, prod(resid_size), N, f, bc, u0_size,
-        tspan, alg.ode_alg)
+        tspan, alg.ode_alg, u0)
 
     ode_fn = (du, u) -> solve_internal_odes!(du, u, prob.p, cur_nshoot, nodes)
     sd_ode = alg.jac_alg.nonbc_diffmode isa AbstractSparseADType ?
@@ -122,7 +117,8 @@ function __solve_nlproblem!(alg::MultipleShooting, bcresid_prototype, u_at_nodes
         ode_fn, similar(u_at_nodes, cur_nshoot * N), u_at_nodes)
 
     bc_fn = (du, u) -> __multiple_shooting_mpoint_loss_bc!(du, u, prob.p,
-        cur_nshoot, nodes, iip, solve_internal_odes!, N, f, bc, u0_size, tspan, alg.ode_alg)
+        cur_nshoot, nodes, iip, solve_internal_odes!, N, f, bc, u0_size, tspan, alg.ode_alg,
+        u0)
     sd_bc = alg.jac_alg.bc_diffmode isa AbstractSparseADType ?
             SymbolicsSparsityDetection() : NoSparsityDetection()
     bc_jac_cache = sparse_jacobian_cache(alg.jac_alg.bc_diffmode,
@@ -144,9 +140,9 @@ function __solve_nlproblem!(alg::MultipleShooting, bcresid_prototype, u_at_nodes
     return nothing
 end
 
-function __multiple_shooting_solve_internal_odes!(resid_nodes, us, p, ::Val{iip}, f,
+function __multiple_shooting_solve_internal_odes!(resid_nodes, us, p, ::Val{iip}, f::F,
     cur_nshoots::Int, nodes, tspan, u0_size, N, alg::MultipleShooting,
-    ensemblealg, kwargs) where {iip}
+    ensemblealg, kwargs) where {iip, F}
     ts_ = Vector{Vector{typeof(first(tspan))}}(undef, cur_nshoots)
     us_ = Vector{Vector{typeof(us)}}(undef, cur_nshoots)
 
@@ -194,7 +190,7 @@ function __multiple_shooting_mpoint_jacobian!(J, us, p, resid_bc, resid_nodes,
 end
 
 @views function __multiple_shooting_2point_loss!(resid, us, p, cur_nshoots::Int, nodes,
-    ::Val{iip}, solve_internal_odes!, resida_len, residb_len, N, bca, bcb) where {iip}
+    ::Val{iip}, solve_internal_odes!::S, resida_len, residb_len, N, bca, bcb) where {iip, S}
     resid_ = resid[(resida_len + 1):(end - residb_len)]
     solve_internal_odes!(resid_, us, p, cur_nshoots, nodes)
 
@@ -216,13 +212,14 @@ end
 end
 
 @views function __multiple_shooting_mpoint_loss_bc!(resid_bc, us, p, cur_nshoots::Int,
-    nodes, ::Val{iip}, solve_internal_odes!, N, f, bc, u0_size, tspan, ode_alg) where {iip}
+    nodes, ::Val{iip}, solve_internal_odes!::S, N, f, bc, u0_size, tspan,
+    ode_alg, u0) where {iip, S}
     _resid_nodes = similar(us, cur_nshoots * N)
 
     # NOTE: We need to recompute this to correctly propagate the dual numbers / gradients
     _us, _ts = solve_internal_odes!(_resid_nodes, us, p, cur_nshoots, nodes)
 
-    odeprob = ODEProblem{iip}(f, reshape(us[1:N], u0_size), tspan, p)
+    odeprob = ODEProblem{iip}(f, u0, tspan, p)
     total_solution = SciMLBase.build_solution(odeprob, ode_alg, _ts, _us)
 
     if iip
@@ -235,14 +232,14 @@ end
 end
 
 @views function __multiple_shooting_mpoint_loss!(resid, us, p, cur_nshoots::Int, nodes,
-    ::Val{iip}, solve_internal_odes!, resid_len, N, f, bc, u0_size, tspan,
-    ode_alg) where {iip}
+    ::Val{iip}, solve_internal_odes!::S, resid_len, N, f, bc, u0_size, tspan,
+    ode_alg, u0) where {iip, S}
     resid_bc = resid[1:resid_len]
     resid_nodes = resid[(resid_len + 1):end]
 
     _us, _ts = solve_internal_odes!(resid_nodes, us, p, cur_nshoots, nodes)
 
-    odeprob = ODEProblem{iip}(f, reshape(us[1:N], u0_size), tspan, p)
+    odeprob = ODEProblem{iip}(f, u0, tspan, p)
     total_solution = SciMLBase.build_solution(odeprob, ode_alg, _ts, _us)
 
     if iip

src/utils.jl

@@ -71,27 +71,30 @@ function __maybe_matmul!(z, A, b, α = eltype(z)(1), β = eltype(z)(0))
 end
 
 ## Easier to dispatch
-eval_bc_residual(pt, bc, sol, p) = eval_bc_residual(pt, bc, sol, p, sol.t)
-eval_bc_residual(_, bc, sol, p, t) = bc(sol, p, t)
-function eval_bc_residual(::TwoPointBVProblem, (bca, bcb), sol, p, t)
+eval_bc_residual(pt, bc::BC, sol, p) where {BC} = eval_bc_residual(pt, bc, sol, p, sol.t)
+eval_bc_residual(_, bc::BC, sol, p, t) where {BC} = bc(sol, p, t)
+function eval_bc_residual(::TwoPointBVProblem, (bca, bcb)::BC, sol, p, t) where {BC}
     ua = sol isa AbstractVector ? sol[1] : sol(first(t))
     ub = sol isa AbstractVector ? sol[end] : sol(last(t))
     resida = bca(ua, p)
     residb = bcb(ub, p)
     return (resida, residb)
 end
 
-eval_bc_residual!(resid, pt, bc!, sol, p) = eval_bc_residual!(resid, pt, bc!, sol, p, sol.t)
-eval_bc_residual!(resid, _, bc!, sol, p, t) = bc!(resid, sol, p, t)
-@views function eval_bc_residual!(resid, ::TwoPointBVProblem, (bca!, bcb!), sol, p, t)
+function eval_bc_residual!(resid, pt, bc!::BC, sol, p) where {BC}
+    return eval_bc_residual!(resid, pt, bc!, sol, p, sol.t)
+end
+eval_bc_residual!(resid, _, bc!::BC, sol, p, t) where {BC} = bc!(resid, sol, p, t)
+@views function eval_bc_residual!(resid, ::TwoPointBVProblem, (bca!, bcb!)::BC, sol, p,
+    t) where {BC}
     ua = sol isa AbstractVector ? sol[1] : sol(first(t))
     ub = sol isa AbstractVector ? sol[end] : sol(last(t))
     bca!(resid.resida, ua, p)
     bcb!(resid.residb, ub, p)
     return resid
 end
-@views function eval_bc_residual!(resid::Tuple, ::TwoPointBVProblem, (bca!, bcb!), sol, p,
-    t)
+@views function eval_bc_residual!(resid::Tuple, ::TwoPointBVProblem, (bca!, bcb!)::BC, sol,
+    p, t) where {BC}
     ua = sol isa AbstractVector ? sol[1] : sol(first(t))
     ub = sol isa AbstractVector ? sol[end] : sol(last(t))
     bca!(resid[1], ua, p)