Merge pull request #1833 from ErikQQY/thread

Add more optional threading
SciML · Jan 11, 2023 · f377a9d · f377a9d
2 parents 64a83eb + a5199fa
commit f377a9d
Show file tree

Hide file tree

Showing 14 changed files with 2,898 additions and 757 deletions.
diff --git a/src/algorithms.jl b/src/algorithms.jl
diff --git a/src/caches/adams_bashforth_moulton_caches.jl b/src/caches/adams_bashforth_moulton_caches.jl
@@ -461,7 +461,7 @@ function alg_cache(alg::VCAB4, u, rate_prototype, ::Type{uEltypeNoUnits},
     rtmp = zero(u)
     ratmp = similar(u, uEltypeNoUnits)
     recursivefill!(ratmp, false)
-    rk4cache = RK4Cache(u, uprev, rk1, rk2, rk3, rk4, rk, rtmp, ratmp)
+    rk4cache = RK4Cache(u, uprev, rk1, rk2, rk3, rk4, rk, rtmp, ratmp, trivial_limiter!, trivial_limiter!, False())
     fsalfirst = zero(rate_prototype)
     k4 = zero(rate_prototype)
     dts = fill(zero(dt), 4)
@@ -556,7 +556,7 @@ function alg_cache(alg::VCAB5, u, rate_prototype, ::Type{uEltypeNoUnits},
     rtmp = zero(u)
     ratmp = similar(u, uEltypeNoUnits)
     recursivefill!(ratmp, false)
-    rk4cache = RK4Cache(u, uprev, rk1, rk2, rk3, rk4, rk, rtmp, ratmp)
+    rk4cache = RK4Cache(u, uprev, rk1, rk2, rk3, rk4, rk, rtmp, ratmp, trivial_limiter!, trivial_limiter!, False())
     fsalfirst = zero(rate_prototype)
     k4 = zero(rate_prototype)
     dts = fill(zero(dt), 5)
@@ -760,7 +760,7 @@ function alg_cache(alg::VCABM4, u, rate_prototype, ::Type{uEltypeNoUnits},
     rtmp = zero(u)
     ratmp = similar(u, uEltypeNoUnits)
     recursivefill!(ratmp, false)
-    rk4cache = RK4Cache(u, uprev, rk1, rk2, rk3, rk4, rk, rtmp, ratmp)
+    rk4cache = RK4Cache(u, uprev, rk1, rk2, rk3, rk4, rk, rtmp, ratmp, trivial_limiter!, trivial_limiter!, False())
     fsalfirst = zero(rate_prototype)
     k4 = zero(rate_prototype)
     dts = fill(zero(dt), 4)
@@ -863,7 +863,7 @@ function alg_cache(alg::VCABM5, u, rate_prototype, ::Type{uEltypeNoUnits},
     rtmp = zero(u)
     ratmp = similar(u, uEltypeNoUnits)
     recursivefill!(ratmp, false)
-    rk4cache = RK4Cache(u, uprev, rk1, rk2, rk3, rk4, rk, rtmp, ratmp)
+    rk4cache = RK4Cache(u, uprev, rk1, rk2, rk3, rk4, rk, rtmp, ratmp, trivial_limiter!, trivial_limiter!, False())
     fsalfirst = zero(rate_prototype)
     k4 = zero(rate_prototype)
     dts = fill(zero(dt), 5)

diff --git a/src/caches/high_order_rk_caches.jl b/src/caches/high_order_rk_caches.jl
@@ -1,4 +1,4 @@
-@cache struct TanYam7Cache{uType, rateType, uNoUnitsType, TabType} <:
+@cache struct TanYam7Cache{uType, rateType, uNoUnitsType, TabType, StageLimiter, StepLimiter, Thread} <:
               OrdinaryDiffEqMutableCache
     u::uType
     uprev::uType
@@ -17,6 +17,9 @@
     atmp::uNoUnitsType
     k::rateType
     tab::TabType
+    stage_limiter!::StageLimiter
+    step_limiter!::StepLimiter
+    thread::Thread
 end
 
 function alg_cache(alg::TanYam7, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -40,7 +43,7 @@ function alg_cache(alg::TanYam7, u, rate_prototype, ::Type{uEltypeNoUnits},
     recursivefill!(atmp, false)
     k = zero(rate_prototype)
     TanYam7Cache(u, uprev, k1, k2, k3, k4, k5, k6, k7, k8, k9, k10, utilde, tmp, atmp, k,
-                 tab)
+                 tab, alg.stage_limiter!, alg.step_limiter!, alg.thread)
 end
 
 function alg_cache(alg::TanYam7, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -50,7 +53,7 @@ function alg_cache(alg::TanYam7, u, rate_prototype, ::Type{uEltypeNoUnits},
     TanYam7ConstantCache(constvalue(uBottomEltypeNoUnits), constvalue(tTypeNoUnits))
 end
 
-@cache struct DP8Cache{uType, rateType, uNoUnitsType, TabType} <: OrdinaryDiffEqMutableCache
+@cache struct DP8Cache{uType, rateType, uNoUnitsType, TabType, StageLimiter, StepLimiter, Thread} <: OrdinaryDiffEqMutableCache
     u::uType
     uprev::uType
     k1::rateType
@@ -81,6 +84,9 @@ end
     tmp::uType
     atmp::uNoUnitsType
     tab::TabType
+    stage_limiter!::StageLimiter
+    step_limiter!::StepLimiter
+    thread::Thread
 end
 
 function alg_cache(alg::DP8, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -121,7 +127,7 @@ function alg_cache(alg::DP8, u, rate_prototype, ::Type{uEltypeNoUnits},
     DP8Cache(u, uprev, k1, k2, k3, k4, k5, k6, k7, k8, k9, k10, k11, k12, k13, k14, k15,
              k16, kupdate,
              udiff, bspl, dense_tmp3, dense_tmp4, dense_tmp5, dense_tmp6, dense_tmp7,
-             utilde, tmp, atmp, tab)
+             utilde, tmp, atmp, tab, alg.stage_limiter!, alg.step_limiter!, alg.thread)
 end
 
 function alg_cache(alg::DP8, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -131,7 +137,7 @@ function alg_cache(alg::DP8, u, rate_prototype, ::Type{uEltypeNoUnits},
     DP8ConstantCache(constvalue(uBottomEltypeNoUnits), constvalue(tTypeNoUnits))
 end
 
-@cache struct TsitPap8Cache{uType, rateType, uNoUnitsType, TabType} <:
+@cache struct TsitPap8Cache{uType, rateType, uNoUnitsType, TabType, StageLimiter, StepLimiter, Thread} <:
               OrdinaryDiffEqMutableCache
     u::uType
     uprev::uType
@@ -153,6 +159,9 @@ end
     atmp::uNoUnitsType
     k::rateType
     tab::TabType
+    stage_limiter!::StageLimiter
+    step_limiter!::StepLimiter
+    thread::Thread
 end
 
 function alg_cache(alg::TsitPap8, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -179,7 +188,7 @@ function alg_cache(alg::TsitPap8, u, rate_prototype, ::Type{uEltypeNoUnits},
     atmp = similar(u, uEltypeNoUnits)
     recursivefill!(atmp, false)
     TsitPap8Cache(u, uprev, k1, k2, k3, k4, k5, k6, k7, k8, k9, k10, k11, k12, k13, utilde,
-                  tmp, atmp, k, tab)
+                  tmp, atmp, k, tab, alg.stage_limiter!, alg.step_limiter!, alg.thread)
 end
 
 function alg_cache(alg::TsitPap8, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -189,7 +198,7 @@ function alg_cache(alg::TsitPap8, u, rate_prototype, ::Type{uEltypeNoUnits},
     TsitPap8ConstantCache(constvalue(uBottomEltypeNoUnits), constvalue(tTypeNoUnits))
 end
 
-@cache struct PFRK87Cache{uType, rateType, uNoUnitsType, TabType} <:
+@cache struct PFRK87Cache{uType, rateType, uNoUnitsType, TabType, StageLimiter, StepLimiter, Thread} <:
               OrdinaryDiffEqMutableCache
     u::uType
     uprev::uType
@@ -211,6 +220,9 @@ end
     atmp::uNoUnitsType
     k::rateType
     tab::TabType
+    stage_limiter!::StageLimiter
+    step_limiter!::StepLimiter
+    thread::Thread
 end
 
 function alg_cache(alg::PFRK87, u, rate_prototype, ::Type{uEltypeNoUnits},
@@ -237,7 +249,7 @@ function alg_cache(alg::PFRK87, u, rate_prototype, ::Type{uEltypeNoUnits},
     atmp = similar(u, uEltypeNoUnits)
     recursivefill!(atmp, false)
     PFRK87Cache(u, uprev, k1, k2, k3, k4, k5, k6, k7, k8, k9, k10, k11, k12, k13, utilde,
-                tmp, atmp, k, tab)
+                tmp, atmp, k, tab, alg.stage_limiter!, alg.step_limiter!, alg.thread)
 end
 
 function alg_cache(alg::PFRK87, u, rate_prototype, ::Type{uEltypeNoUnits},