Dependency Graph Framework for Automatic Tendency Ordering and Caching

[nefrathenrici]

Purpose

This SDI describes a system to replace manual orchestration of tendency computations with an automatic, declarative system.

DependencyGraph constructs dependency graphs for computing tendencies from prognostic variables, topologically sorts them, and allows configurable evaluation of nodes with caching.

Y (prognostic) → computed vars → Yₜ (tendencies)

Core Strategy

ClimaAtmos computes tendencies through the chain of precomputed quantities. These are manually ordered, all cached in global memory, and recomputed even when not needed.

VarManager provides a declarative dependency system to redefine how tendencies are computed.

1. Users define what each variable depends on using var_dependencies and tendency_dependencies:

var_dependencies(::AtmosModel, ::FieldName{(:c, :J)}) = (@name(c.uₕ),)
tendency_dependencies(::AtmosModel, ::FieldName{(:c, :ρ)}) = (@name(f.ρ), @name(f.u³))

2. Users define how to compute vars and tendencies:

compute_var(::AtmosModel, ::FieldName{(:c, :J)}, vars, t) =
    Fields.local_geometry_field(vars[@name(c.uₕ)]).J
compute_tendency(::AtmosModel, ::FieldName{(:c, :ρ)}, vars, t) =
    @. lazy(-(ᶜdivᵥ(vars[@name(f.ρ)] * vars[@name(f.u³)])))

3. This is used to build an acyclic dependency graph that is sorted to obtain the evaluation order for computed variables.

build_dependency_graph(tend_names, prog_names, model)
get_evaluation_order(graph)

4. The graph (or a specific tendency) can be evaluated using a specific caching strategy. The initial strategies will be to cache everything or lazify everything.

evaluate_graph!(Yₜ, Y, graph, model, t, EagerGlobalCaching())

Benefits

• Automatic ordering: Topological graph sort ensures correct and consistent evaluation order
• Configurable caching: Choose what to cache vs. keep lazy. This will allow us to separate performance concerns from ClimaAtmos code.
• Error tracing: Failures report the full dependency chain. The dependency graph will be easily to inspect.
• Kernel fusion: Independent computations in the tree can be fused automatically in the future.

Risks

• Ensuring this system is performant, especially when using eager caching
• This may be difficult to switch to and may require more boilerplate code.
• Requiring acyclic graph may limit timestepping schemes. For example, how will this work with callbacks and implicit/explicit splitting?
• ClimaAtmos can not be changed incrementally to use this, will need full overhaul.

Producers

@nefrathenrici @dennisYatunin

Components

• DependencyGraph: DAG of variable dependencies
• VarCache: Runtime storage for computed variables
• VarCachingStrategy: Abstract strategy for caching policies
• Debugging and graph inspecting utilities

Inputs

A description of the inputs to the solution (designs, references, equations, closures, discussions etc).

Results and deliverables

A description of the key results, deliverables, quality expectations, and performance metrics.

Task breakdown

A preliminary list of PRs and a preliminary timeline of PRs, milestones, and key results.

• Task 1
• Task 2
  ...

Reviewers

@dennisYatunin

[petebachant]

Cool idea. Is this to automatically fix any of the suboptimal intermediate variable creation for the Jacobian update, or has it historically been easy to define these in the wrong order, or maybe both?

[nefrathenrici]

Cool idea. Is this to automatically fix any of the suboptimal intermediate variable creation for the Jacobian update, or has it historically been easy to define these in the wrong order, or maybe both?

This will definitely be used to eliminate unnecessary intermediate variables, but I'd say the primary goal is to isolate variable caching and tendency evaluation from the model/tendency code.