ClimaTimeSteppers currently manages the cache when it could be cache-agnostic.
Current calls within CTS:
- cache! before explicit tendency (imex_ark.jl:176) full refresh before T_exp_T_lim!
- cache_imp! before/within Newton (imex_ark.jl:186-188, :259) implicit refresh before T_imp! and Wfact
- cache! at step end (imex_ark.jl:84) leave p.precomputed consistent with final u for callbacks/diagnostics
- f.cache! at init (integrators.jl:209 -> functions.jl:94) initial cache populated at t0
Possible solution within Atmos:
- Cache refresh within explicit tendency: Call
set_precomputed_quantities! at the start of remaining_tendency!.
- Implicit cache refresh within implicit callers: Call
set_implicit_precomputed_quantities! at the start of:
- implicit_tendency!
- update_jacobian!
- Callback that calls
set_precomputed_quantities! at the end of each step for callbacks so cache-reading callbacks/diagnostics see a consistent cache. The callback can use a flag in the cache to refresh the cache on-demand only if a callback needs it -- see Performance concerns below for more info.
- Cache is already prepared via
build_cache -> set_precomputed_quantities! in cache.jl:186
Performance concerns
First stage redundantly does a full refresh every step boundary: CTS currently guards against refreshing the cache in the first stage, since the state at the end of previous step is the same as the start of the current step. This redundancy is made worse if we also refresh the cache at the step end for callbacks.
Possible Mitigation: New function that refreshes the cache on-demand for callbacks with a cache flag. The flag uses integrator.step, which is bumped once at the top of each step. The first cache-reading callback after a step finds it stale and refreshes once at the final u while later readers skip it. This assumes that no callbacks modify the cache.
function lazy_set_precomputed!(integrator)
(; u, p, t, step) = integrator
if p.precomputed_synced_step[] != step
set_precomputed_quantities!(u, p, t)
p.precomputed_synced_step[] = step
end
return nothing
end
Within the Newton solve, T_imp! and Wfact each refresh the implicit quantities at the same iterate, one redundant implicit-only refresh per iteration. With the default max_iters = 1 that's ~1 extra refresh per implicit stage (e.g. ~3/step for ARS343).
ClimaTimeSteppers currently manages the cache when it could be cache-agnostic.
Current calls within CTS:
Possible solution within Atmos:
set_precomputed_quantities!at the start ofremaining_tendency!.set_implicit_precomputed_quantities!at the start of:set_precomputed_quantities!at the end of each step for callbacks so cache-reading callbacks/diagnostics see a consistent cache. The callback can use a flag in the cache to refresh the cache on-demand only if a callback needs it -- see Performance concerns below for more info.build_cache->set_precomputed_quantities!in cache.jl:186Performance concerns
First stage redundantly does a full refresh every step boundary: CTS currently guards against refreshing the cache in the first stage, since the state at the end of previous step is the same as the start of the current step. This redundancy is made worse if we also refresh the cache at the step end for callbacks.
Possible Mitigation: New function that refreshes the cache on-demand for callbacks with a cache flag. The flag uses
integrator.step, which is bumped once at the top of each step. The first cache-reading callback after a step finds it stale and refreshes once at the final u while later readers skip it. This assumes that no callbacks modify the cache.Within the Newton solve, T_imp! and Wfact each refresh the implicit quantities at the same iterate, one redundant implicit-only refresh per iteration. With the default max_iters = 1 that's ~1 extra refresh per implicit stage (e.g. ~3/step for ARS343).