Skip to content

Decouple cache from timestepper #4511

Description

@nefrathenrici

ClimaTimeSteppers currently manages the cache when it could be cache-agnostic.

Current calls within CTS:

  1. cache! before explicit tendency (imex_ark.jl:176) full refresh before T_exp_T_lim!
  2. cache_imp! before/within Newton (imex_ark.jl:186-188, :259) implicit refresh before T_imp! and Wfact
  3. cache! at step end (imex_ark.jl:84) leave p.precomputed consistent with final u for callbacks/diagnostics
  4. f.cache! at init (integrators.jl:209 -> functions.jl:94) initial cache populated at t0

Possible solution within Atmos:

  1. Cache refresh within explicit tendency: Call set_precomputed_quantities! at the start of remaining_tendency!.
  2. Implicit cache refresh within implicit callers: Call set_implicit_precomputed_quantities! at the start of:
    • implicit_tendency!
    • update_jacobian!
  3. 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.
  4. 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).

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions