Tracers rework

[FlorianDeconinck]

As part of the work to up-skill the dynamics (GEOS-ESM/PyFV3#4) we need to be able to advect N tracers. This echoes a larger issue where the tracers in dynamics have been fixed (with downstream effects on both pyFV3 and NDSL) as a dictionary of quantities, but physics are currently implementing some as a 4 dimensional array where the tracers are the 4th D.

Two goals emerge for the rework

• for dynamics (pyFV3): we need to advect N tracers
• for physics: we need to discussing improving upon the current Fortran strategy of indexing.

Therefore we propose the following plan:

• Part 1: rework the dynamics tracers to allow N Quantities to be advected. Keep the ongoing strategy but remove every hard-coding and extend to the capacity to describe N tracers, including nameless tracers.
• Part 2: Design an unified in and off stencil strategy for Tracers, but larger, for any bundle which isn't really a "dimension" but rather a list
• Part 3: Implement the results of Part 2 with physics as the base case and retrofit them into dynamics

Timeline:

• Part 1: End of 2024
• Part 2 & Part 3: TBD

Below we keep notes for the re-design.

Authors: @oelbert @twicki @FlorianDeconinck

[FlorianDeconinck]

Design for Bundling finite elements in and off stencils

We will run with the Tracers idea here, but this is re-usable for any finite list of items that is mapped on the grid.

The dual issue

Right now we have dynamics implementing tracers with an off-stencil solution e.g.:

@gtscript.stencil(...)
def stencil_operating_on_q(q: FloatField, ...):
  with computation(PARALLEL), interval(...):
    ... # preamble calculation
     ... q ...

@gtscript.stencil(...) 
def stencil_operating_on_ice_and_snow(qsnow: FloatField, qice: FloatField, ...):
  with computation(PARALLEL), interval(...):
     qsnow = ... qice ...

 def solver(tracers: dict):
   for tracer in tracers:
     stencil_operating_on_q(tracer)
   stencil_operating_on_ice_and_snow(tracers["snow"], tracers["ice"]

➕

• Very clear and clean coding: separation of concerns, string named tracers for call

➖

• Bad locality for any preamble calculation (see example)
• If code updating tracer is a small amount of code in a larger stencil, need to either buffer the update for later (if possible) or break the code into mulitple stencils

In physics (PBL, UW, Aerosol Activation, ...) we have a lot of need for in-stencil computation which are done using a 4D array (with data dimensions), e.g.:

@gtscript.stencil(...)
def physics_stencil(q: FloatField4D, n_tracers:int, index_for_ice: int ...):
  with computation(PARALLEL), interval(...):
    ... # preamble calculation
    while i_tracer < n_tracers:
      q[0, 0, 0][i_tracer] = ...
   ...
   q[0, 0, 0][index_for_ice] = ...

➕

• Good locality for performance
• Able to insert tracer code in bigger solvers easily

➖

• More difficult to analyze performance (due to the while loop) and potential performance lock (impossible to break out in multiple kernel)
• Less clear when dealing with "generic tracer update code"
• Use of global while concept which is, in fact, a for
• Relies on int indexing: bad software practice, poor readability

Some embryos of solutions

It seems both are valid use cases. We would need to work torward unifying the frontend for both, which would require gt4py changes.

First option is to allow in-stencil for looping + name index testing, e.g.

@gtscript.stencil(...)
def stencil(tracers: FloatFieldList):
  with computation(PARALLEL), interval(...):
    for q in tracers:
       if q.name == "tke":
            q = ...
        else:
            q = ...

Another form of the above, with a more verbose approach that could hide some of the implementing detail a little bit more. Requires element indexing via string in stencil and custom iterators

@gtscript.stencil(...)
def stencil(tracers: FloatFieldList):
  with computation(PARALLEL), interval(...):
    for q in tracers.exclude(["tke"]):
            q = ...
    tracers["tke"] = ...

Both those try to combine pros and leaves cons aside, allowing performance tuning while retaining off-stencil capacity and verbosity.

[FlorianDeconinck]

Ping @bensonr @fmalatino

This the beginning of a larger discussion to be had in the new year to evolve the concepts around list of float fields (like tracers) to match both in and off stencil usage.

[FlorianDeconinck]

PR for Part 1: NOAA-GFDL/PyFV3#37

[FlorianDeconinck]

An issue arise with the Part1 with orchestration. In orchestration, the memory needs to be staticly defined at JIT times, which means dict and list are an issue. So far, the system allows for one deviation: unrolling.

If you write

for key, value in my_dict.items():
    if key in exlude_list:
        pass
    else:
       stencil(value)

You get the code that you expect but the for was unrolled, e.g. the code was repeated for each element of my_dict that are not skipped. Compile times are fine for 7 tracers, not so fine on 200+.

The system allows you to specify dace.nounroll on range/map but it doesn't apply to dict. I tried a few options by redefining a custom object, a 4D array, but nothing is both understandable (read pythonic and expected) and works.

[bensonr]

@FlorianDeconinck - I'm late to the discussion as I try to recover my inbox from time off at the end of the year. Knowing we are reading in the tracers as yaml key-value pairs, how are we going to define the tracers that fall into the various classes that have been maintained within the original FV3 - namely hydrometeor tracers (vapor, water, frozen, etc.), chemistry tracers, non-advected (prognostic and diagnostic), as well as number quantities?

[FlorianDeconinck]

The idea is to have named tracers when those make sense (so the hydrometeor are all named because in the dynamics they are indeed called by name, e.g. tracers["vapor"]) and we carry relevant list of names for every special criteria we need in the code . E.g. "advected" will be a static list of names that will be read in to figure out which tracers are advected.

We use the JIT nature of the system to go from a index-only system that is in the current Fortran, to a verbose string-based systen. At runtime all strings have been flushed out and we are effectively executing index list. But for the users things can be query by name for ease of use

[bensonr]

I'd like to not rely on named tracers for behavior if we can instead use the extended attributes (e.g hydrometeor and phase) to determine how the tracers are used for the various tracer schemes. For example the computations where we determine liquid mass and apply the specific heat to that mass as opposed to vapor or frozen (solid).

[FlorianDeconinck]

Well the current solution is based on the need of the implementation of the current algorithm, which rely on pre-set indices into a 4D array, which is akin to named tracers, with the verbosity of the string replaced by the more opaque index. As always, I don't alter algorithmic choices - I just try to upgrade what i can to more modern standards.

That said, what you describe is totally do-able. The main limitation is that we need a system that freezes any indexation/retrival of tracers when DaCe JIT is called, since DaCe demands to know exactly the shape of the memory to JIT properly. That system will be similar I expect between bluntly named tracers, and more finely define metadata on the tracers that allow you to write something like

for q in tracers.hydrometeors:
   ...